ezafe_xlmroberta.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HmnQcelhsqyx"
   },
   "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": "2nkF081WGC2u"
   },
   "outputs": [],
   "source": [
    "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 XLMRobertaForTokenClassification, XLMRobertaTokenizer\n",
    "\n",
    "\n",
    "random.seed(17)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\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/roberta_model/'\n",
    "\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 = 2"
   ]
  },
  {
   "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_ezafe_tags, ezafe_tags = [], [TAG_PAD_OR_MASK_TOKEN]\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",
    "                ezafe_tags.extend([int(ezafe_tag)] + [TAG_PAD_OR_MASK_TOKEN for _ in range(len(word) - 1)])\n",
    "            else:\n",
    "                assert len(sent) == len(ezafe_tags)\n",
    "                sents.append(sent + [END_TOKEN])\n",
    "                all_ezafe_tags.append(ezafe_tags + [TAG_PAD_OR_MASK_TOKEN])\n",
    "                \n",
    "                sent = [BEGIN_TOKEN]\n",
    "                ezafe_tags = [TAG_PAD_OR_MASK_TOKEN]\n",
    "\n",
    "    return sents, all_ezafe_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 EzafeDataset(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": [
    "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",
    "try:\n",
    "    model = XLMRobertaForTokenClassification.from_pretrained(model_dir + '_' + str(last_i))\n",
    "    print('model is loaded')\n",
    "except:\n",
    "    model = XLMRobertaForTokenClassification.from_pretrained('xlm-roberta-base')\n",
    "\n",
    "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",
    "train_dataset = EzafeDataset(train_data)\n",
    "train_sampler = MySampler(train_dataset, last_i)\n",
    "\n",
    "valid_dataset = EzafeDataset(valid_data)\n",
    "\n",
    "print(len(train_dataset))\n",
    "print(len(valid_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 231
    },
    "colab_type": "code",
    "id": "SrptV6bQC7_e",
    "outputId": "ea881449-84d5-4360-f51d-af5a96a4cf31"
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "ignored",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-00adea232d82>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m train_data_loader = DataLoader(dataset=train_dataset,                                                         \n\u001b[0m\u001b[1;32m      4\u001b[0m                                \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m                                \u001b[0mcollate_fn\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcollate_fn\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'DataLoader' is not defined"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "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",
    "model.to(device)\n",
    "\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    for i, (x, y_true) in enumerate(train_data_loader):\n",
    "        model.train()\n",
    "        i += last_i\n",
    "\n",
    "        if i % 1000 == 0 and i != 0:\n",
    "            print(i)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\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",
    "        loss = criterion(y_pred.view(-1, y_pred.shape[-1]), y_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('validating...')\n",
    "            mode.evaluate()\n",
    "            true_labels, valid_pred_labels = [], []\n",
    "            with torch.no_grad():\n",
    "                for x, y_true in valid_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",
    "                    valid_true_labels.extend(y_true.reshape(-1).cpu().numpy().tolist())\n",
    "                    valid_pred_labels.extend(torch.argmax(y_pred, -1).reshape(-1).cpu().numpy().tolist())\n",
    "\n",
    "            valid_true_labels_, valid_pred_labels_ = np.array(valid_true_labels), np.array(valid_pred_labels)\n",
    "            valid_pad_or_mask_ids = np.where(valid_true_labels_ == TAG_PAD_OR_MASK_TOKEN)[0]\n",
    "            valid_true_labels_ = np.delete(valid_true_labels_, valid_pad_or_mask_ids)\n",
    "            valid_pred_labels_ = np.delete(valid_pred_labels_, valid_pad_or_mask_ids)\n",
    "  \n",
    "            valid_f1 = classification_report(valid_true_labels_, valid_pred_labels_, digits=4)\n",
    "              \n",
    "            summary = f'epoch: {epoch + 1} | step: {i}:\\n\\n{valid_f1}\\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",
    "            model.save_pretrained(model_dir + str(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7z1vgLV57abh"
   },
   "outputs": [],
   "source": [
    "# testing\n",
    "test_dataset = EzafeDataset(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 = 22500\n",
    "model = XLMRobertaForTokenClassification.from_pretrained(model_dir + str(i))\n",
    "\n",
    "model.to(device)\n",
    "model.evaluate()\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": "ezafe_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
}