try7_XLMR.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "CUDA_VISIBLE_DEVICES=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data1/debajyoti/colie/colenv/lib/python3.8/site-packages/tqdm/auto.py:21: 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 torch\n",
    "import torch.nn as nn\n",
    "import os\n",
    "import logging\n",
    "import numpy as np\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "import pandas as pd\n",
    "\n",
    "from transformers import LongformerModel, AutoTokenizer, LongformerForSequenceClassification, LongformerForMultipleChoice\n",
    "from transformers import AutoTokenizer, XLMRobertaModel\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import f1_score, accuracy_score, classification_report\n",
    "\n",
    "logging.basicConfig(filename=f'./logs/train_{time.asctime().replace(\" \",\"_\")}.log', filemode='w', level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\n",
    "\n",
    "# Create a logger object\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "\n",
    "# Create a stream handler to print log messages to the console\n",
    "console_handler = logging.StreamHandler()\n",
    "console_handler.setLevel(logging.INFO)\n",
    "formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')\n",
    "console_handler.setFormatter(formatter)\n",
    "logger.addHandler(console_handler)\n",
    "\n",
    "torch.manual_seed(40)\n",
    "np.random.seed(40)\n",
    "random.seed(40)\n",
    "torch.cuda.manual_seed(40)\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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>BOOK_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>7616_1.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>7616_2.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>7616_3.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7616_4.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>7616_5.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143025</th>\n",
       "      <td>5677_92.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143026</th>\n",
       "      <td>5677_93.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143027</th>\n",
       "      <td>5677_94.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143028</th>\n",
       "      <td>5677_95.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143029</th>\n",
       "      <td>5677_96.txt</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>143030 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            BOOK_id\n",
       "0        7616_1.txt\n",
       "1        7616_2.txt\n",
       "2        7616_3.txt\n",
       "3        7616_4.txt\n",
       "4        7616_5.txt\n",
       "...             ...\n",
       "143025  5677_92.txt\n",
       "143026  5677_93.txt\n",
       "143027  5677_94.txt\n",
       "143028  5677_95.txt\n",
       "143029  5677_96.txt\n",
       "\n",
       "[143030 rows x 1 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Define the path to the CSV file\n",
    "train_csv_file = \"/data1/debajyoti/colie/train.csv\"\n",
    "val_csv_file = \"/data1/debajyoti/colie/valid.csv\"\n",
    "test_csv_file = \"/data1/debajyoti/colie/test.csv\"\n",
    "\n",
    "# Read the CSV file\n",
    "train_labels = pd.read_csv(train_csv_file)\n",
    "val_labels = pd.read_csv(val_csv_file)\n",
    "test_labels = pd.read_csv(test_csv_file)\n",
    "test_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'27993_1.txt'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_labels.BOOK_id[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                text      label\n",
      "0  rifle; Ivan's was a double-barrelled shot-gun ...  Viktorian\n",
      "1  upon the track of the bear. After following it...  Viktorian\n",
      "2  to pull him out with their hands--even had the...  Viktorian\n",
      "3  a slight sparkle of scientific conceit, \"this ...  Viktorian\n",
      "4  bears with a white ring round their necks? Yes...  Viktorian                                                 text      label\n",
      "0  kind good morning, and returned her hearty emb...  Viktorian\n",
      "1  sky, and of the moon, which clothed the old pi...  Viktorian\n",
      "2  left Rome for Augsburg, my mind being much exc...  Viktorian\n",
      "3  thoughts some of the old melodies he knew by h...  Viktorian\n",
      "4  \"But,\" said Henry, \"is it not possible that th...  Viktorian                                                 text\n",
      "0  \"Alas, poor girl!\" said I, \"I fear that her ha...\n",
      "1  to divide her attention between the said garco...\n",
      "2  visitor's disposition to gallantry. However, s...\n",
      "3  says Juvenal, \"'Mors sola fatetur Quantula sin...\n",
      "4  him out in that back passage; the outer door i...\n",
      "(546210, 2) (36257, 2) (143030, 1)\n"
     ]
    }
   ],
   "source": [
    "# Define the path to the train folder\n",
    "train_folder = \"/data1/debajyoti/colie/train/train/\"\n",
    "# Define the path to the validation folder\n",
    "val_folder = \"/data1/debajyoti/colie/valid/valid/\"\n",
    "# Define the path to the test folder\n",
    "test_folder = \"/data1/debajyoti/colie/test/test/\"\n",
    "\n",
    "\n",
    "\n",
    "def create_df(folder, label):\n",
    "    # Initialize empty lists to store the data\n",
    "    text_data = []\n",
    "    labels = []\n",
    "    for index in label.index:\n",
    "        # filename = df_labels.BOOK_id[index]\n",
    "        # print(filename)\n",
    "        # print(df_labels['BOOK_id'][index], df_labels['Epoch'][index])\n",
    "        file_name = label['BOOK_id'][index]  # Assuming 'File Name' is the column name for the file names in the CSV\n",
    "\n",
    "        # Construct the file path\n",
    "        file_path = os.path.join(folder, file_name)\n",
    "\n",
    "        # Read the text from the file\n",
    "        with open(file_path, 'r', encoding='ISO-8859-1') as file:\n",
    "            text = file.read()\n",
    "\n",
    "        # Append the text and label to the respective lists\n",
    "        text_data.append(text)\n",
    "        labels.append(label['Epoch'][index].strip())  # Assuming 'Label' is the column name for the labels in the CSV\n",
    "        # break\n",
    "    return text_data, labels\n",
    "\n",
    "def create_df_test(folder, label):\n",
    "    # Initialize empty lists to store the data\n",
    "    text_data = []\n",
    "    # labels = []\n",
    "    for index in label.index:\n",
    "        # filename = df_labels.BOOK_id[index]\n",
    "        # print(filename)\n",
    "        # print(df_labels['BOOK_id'][index], df_labels['Epoch'][index])\n",
    "        file_name = label['BOOK_id'][index]  # Assuming 'File Name' is the column name for the file names in the CSV\n",
    "\n",
    "        # Construct the file path\n",
    "        file_path = os.path.join(folder, file_name)\n",
    "\n",
    "        # Read the text from the file\n",
    "        with open(file_path, 'r', encoding='ISO-8859-1') as file:\n",
    "            text = file.read()\n",
    "\n",
    "        # Append the text and label to the respective lists\n",
    "        text_data.append(text)\n",
    "        # labels.append(label['Epoch'][index].strip())  # Assuming 'Label' is the column name for the labels in the CSV\n",
    "        # break\n",
    "    return text_data\n",
    "\n",
    "train_data, train_label = create_df(train_folder, train_labels)\n",
    "val_data, val_label = create_df(val_folder, val_labels)\n",
    "test_data = create_df_test(test_folder, test_labels)\n",
    "\n",
    "# Create a dataframe from the lists\n",
    "train = pd.DataFrame({'text': train_data, 'label': train_label})\n",
    "val = pd.DataFrame({'text': val_data, 'label': val_label})\n",
    "test = pd.DataFrame({'text': test_data})\n",
    "print(train.head(), val.head(), test.head())\n",
    "print(train.shape, val.shape, test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_dic = {'Romanticism':0,\n",
    "            'Viktorian':1,\n",
    "            'Modernism':2,\n",
    "            'PostModernism':3,\n",
    "            'OurDays':4}\n",
    "train['label'] = train['label'].map(label_dic)\n",
    "val['label'] = val['label'].map(label_dic)\n",
    "# test['label'] = test['label'].map(label_dic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "483268    1128\n",
      "483267    1068\n",
      "521384    1065\n",
      "483265    1034\n",
      "81542     1020\n",
      "          ... \n",
      "470405       1\n",
      "130188       1\n",
      "217335       1\n",
      "351867       1\n",
      "368135       1\n",
      "Name: text, Length: 546210, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# Length of text\n",
    "def length (txt):\n",
    "    length = len(txt.split())\n",
    "    return length\n",
    "\n",
    "txt_length = train['text'].apply(lambda x: length(x))\n",
    "print(txt_length.sort_values(ascending = False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "label\n",
       "1    16938\n",
       "2    14848\n",
       "3     1713\n",
       "4     1600\n",
       "0     1158\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val['label'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = LongformerModel.from_pretrained(\"allenai/longformer-base-4096\")\n",
    "# tokenizer = AutoTokenizer.from_pretrained(\"allenai/longformer-base-4096\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"xlm-roberta-base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/2732 [00:00<?, ?it/s]"
     ]
    }
   ],
   "source": [
    "max_length= 500\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, tokenizer, df):\n",
    "        # Initialize thetokenizer\n",
    "        self.tokenizer = tokenizer\n",
    "\n",
    "        self.df = df\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        # Get the text and label from the dataframe\n",
    "        text = self.df.iloc[index]['text']\n",
    "        label = self.df.iloc[index]['label']\n",
    "\n",
    "        # Tokenize the text and convert it to input IDs\n",
    "        inputs = self.tokenizer(\n",
    "            text,\n",
    "            None,\n",
    "            add_special_tokens=False,\n",
    "            max_length=max_length,\n",
    "            padding='max_length',\n",
    "            truncation=True,\n",
    "            return_token_type_ids=False,\n",
    "            return_attention_mask=True,\n",
    "            return_tensors='pt',\n",
    "        )\n",
    "\n",
    "\n",
    "        # Return the input IDs and label as PyTorch tensors\n",
    "        return {\n",
    "            'input_ids': inputs['input_ids'][0],\n",
    "            'attention_mask': inputs['attention_mask'][0],\n",
    "            # 'token_type_ids': inputs['token_type_ids'][0],\n",
    "            'label': torch.tensor(label, dtype=torch.int64),\n",
    "        }\n",
    "\n",
    "# datasetclass = CustomDataset(tokenizer, train)\n",
    "train_dataset = CustomDataset(tokenizer, train)\n",
    "val_dataset = CustomDataset(tokenizer, val)\n",
    "# test_dataset = CustomDataset(tokenizer, test)\n",
    "\n",
    "# DataLoader\n",
    "batch_size = 200\n",
    "train_dataloader = tqdm(DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=64))\n",
    "val_dataloader = tqdm(DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=64))\n",
    "# test_dataloader = tqdm(DataLoader(test_dataset, batch_size=batch_size, shuffle=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerModel(nn.Module):\n",
    "    def __init__(self, num_labels):\n",
    "        super(TransformerModel, self).__init__()\n",
    "        \n",
    "        self.xlmr = XLMRobertaModel.from_pretrained(\"xlm-roberta-base\")\n",
    "        # self.Longformer = LongformerModel.from_pretrained(\"allenai/longformer-base-4096\")\n",
    "        # Freeze all layers except the top 2\n",
    "        for param in self.xlmr.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "        # Unfreeze the parameters of the top 2 layers\n",
    "        for param in self.xlmr.encoder.layer[-3:].parameters():\n",
    "            param.requires_grad = True\n",
    "        \n",
    "        # self.xlnet.resize_token_embeddings(num_tokens)\n",
    "        # self.transformer_encoder = TransformerEncoder(TransformerEncoderLayer(d_model=hidden_size, nhead=num_heads), num_layers=num_layers)\n",
    "        #self.transformer_decoder = TransformerDecoder(TransformerDecoderLayer(d_model=hidden_size, nhead=num_heads), num_layers=num_layers)\n",
    "        #self.transformer = Transformer(nhead=16, num_encoder_layers=6, num_decoder_layers = 6)\n",
    "        self.decoder = nn.Linear(self.xlmr.config.hidden_size, num_labels) \n",
    "        # self.fc1 = nn.Linear(num_tokens, 2)\n",
    "        # self.fc2 = nn.Linear(num_tokens, 2)\n",
    "        # self.fc3 = nn.Linear(num_tokens, 5)\n",
    "        # self.num_classes = num_classes\n",
    "        # self.classifiers = nn.ModuleList([nn.Linear(self.roberta.config.hidden_size, num_classes[i]) for i in range(len(num_classes))])\n",
    "        # self.classifiers = nn.ModuleList([nn.Linear(num_tokens, num_classes[i]) for i in range(len(num_classes))])\n",
    "        # self.tanh = nn.Tanh()\n",
    "\n",
    "    def forward(self, input_ids, attention_mask):  # src = [bsz, seq_len]\n",
    "        long_output = self.xlmr(input_ids=input_ids).pooler_output\n",
    "        # print(long_output.shape)\n",
    "        # roberta_outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)\n",
    "        # last_hidden_state = outputs.last_hidden_state # Shape: (batch_size, sequence_len, hidden_size)\n",
    "        # src_embedded = last_hidden_state\n",
    "        # src_embedded = self.roberta.embeddings(src) # Use RoBERTa model to embed source sequence output: [bsz, seq_len, features,i.e. hidden_dim] [20, 100, 768]\n",
    "        # print(\"shape of roberta embeddings:\", src_embedded.shape)\n",
    "        #tgt_embedded = self.roberta.embeddings(tgt) # Use RoBERTa model to embed target sequence\n",
    "        # src_embedded = src_embedded # output: [bsz, seq_len, features] \n",
    "        # src_embedded = torch.cat([t1,t2,t3, src_embedded],1)\n",
    "\n",
    "        # t1 = torch.cat(src_embedded.size(0) * [t1])\n",
    "        # t2 = torch.cat(src_embedded.size(0) * [t2])\n",
    "        # t3 = torch.cat(src_embedded.size(0) * [t3])\n",
    "        # t = torch.stack([t1,t2,t3], dim=1)\n",
    "        # task_embedded = torch.cat([t, src_embedded],1)  # output shape: [bsz, seq_len, features] [8, 203, 768]\n",
    "\n",
    "        # memory = self.transformer_encoder(src_embedded)  # output shape: [bsz, seq_len, features] [8, 203, 768]\n",
    "        # print(\"shape after transformer encoder layer:\", memory.shape)\n",
    "        #output = self.transformer_decoder(tgt_embedded, memory)\n",
    "        #print(\"shape after transformer decoder layer:\", output.shape)\n",
    "\n",
    "        output = self.decoder(long_output)  # output shape: [bsz, seq_len, vocab_size] [8, 203, 50k]\n",
    "        # print(\"shape after transformer decoder layer:\", output.shape, output.dtype)\n",
    "        # task1_output = self.fc1(output[:,0,:])\n",
    "        # task2_output = self.fc2(output[:,1,:])\n",
    "        # task3_output = self.fc3(output[:,2,:num_classes])\n",
    "        # ae_output = output[:,len(self.num_classes):,:]\n",
    "        # ae_output = output[:,:,:]\n",
    "        # print(\"shape after final linear layer:\", output.shape)\n",
    "        # task_logits = [classifier(pooled_output) for classifier in self.classifiers]\n",
    "        # task_logits = []\n",
    "\n",
    "        # pooled_outputs = [output[:,i,:] for i in range(len(self.num_classes))] # output shape : [bsz, 1, vocab_size]\n",
    "\n",
    "        # for classifier, pooled_output in zip(self.classifiers, pooled_outputs):\n",
    "        #     # pooled_output = self.tanh(pooled_output)\n",
    "        #     logits = classifier(pooled_output)\n",
    "        #     task_logits.append(logits)\n",
    "        \n",
    "        return output\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at xlm-roberta-base were not used when initializing XLMRobertaModel: ['lm_head.layer_norm.weight', 'lm_head.dense.bias', 'lm_head.dense.weight', 'lm_head.bias', 'lm_head.layer_norm.bias']\n",
      "- This IS expected if you are initializing XLMRobertaModel 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 XLMRobertaModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "num_labels = 5\n",
    "\n",
    "model = TransformerModel(num_labels).to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# num_epochs = 5\n",
    "learning_rate = 2e-4\n",
    "class_weights = torch.tensor([12.0, 1.0, 1.0, 8.0, 15.0]).to(device)\n",
    "\n",
    "# Set optimizer and learning rate scheduler\n",
    "criterion = nn.CrossEntropyLoss(weight=class_weights).to(device)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_labels(logit, targets):\n",
    "    \"\"\"\n",
    "    Calculate accuracy and macro F1 score for each class\n",
    "    \"\"\"\n",
    "    # pos = list(task_dict.keys()).index(task_name)\n",
    "    # mask = torch.arange(targets.shape[0]).to(device)\n",
    "    # task_idx = mask[targets[:,pos] != 99]\n",
    "    output = logit\n",
    "    true_label = targets\n",
    "    # print(\"shapes for label:\", output.shape, true_label.shape)\n",
    "    pred_label = torch.argmax(output, dim=1).flatten().tolist()\n",
    "    true_label = true_label.flatten().tolist()\n",
    "\n",
    "\n",
    "    return pred_label, true_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_train_loss = []\n",
    "\n",
    "def train(model: nn.Module) -> None:\n",
    "    model.train()  # turn on train mode\n",
    "    total_loss = 0.\n",
    "    log_interval = 1\n",
    "    start_time = time.time()\n",
    "    num_batches = len(train_dataset) // batch_size\n",
    "    for batch, i in enumerate(train_dataloader):\n",
    "        data, mask, targets = i.values()\n",
    "        data = data.to(device)\n",
    "        mask = mask.to(device)\n",
    "        targets = targets.to(device)\n",
    "        # print(data.dtype)        \n",
    "        # print(data.shape)\n",
    "        # task_logits, ae_output = model(data)\n",
    "        output = model(data, mask)\n",
    "        # t1_out, t2_out, t3_out, auto_output = model(data, t1, t2, t3)\n",
    "        # loss = custom_loss(logits_task1, logits_task2, logits_task3, targets)\n",
    "        # print(\"shape:\", data.shape, targets.flatten().shape)\n",
    "        # print(\"outputtype:\", output.dtype, targets.flatten().dtype)\n",
    "        # print(output)\n",
    "        # targets = targets.float()\n",
    "        loss = criterion(output, targets.flatten())\n",
    "\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)\n",
    "        optimizer.step()\n",
    "\n",
    "        total_loss += loss.item()\n",
    "        if batch % log_interval == 0 and batch > 0:\n",
    "            lr = scheduler.get_last_lr()[0]\n",
    "            ms_per_batch = (time.time() - start_time) * 1000 / log_interval\n",
    "            cur_loss = total_loss / log_interval\n",
    "            # ppl = np.exp(cur_loss)\n",
    "            # print(f'| epoch {epoch:3d} | {batch:5d}/{num_batches:5d} batches | '\n",
    "            #         f'lr {lr:02.7f} | ms/batch {ms_per_batch:5.2f} | '\n",
    "            #         f'loss {cur_loss:5.5f}')\n",
    "            total_loss = 0\n",
    "            start_time = time.time()\n",
    "        \n",
    "        if batch == 100:\n",
    "            break\n",
    "    current_train_loss.append(cur_loss)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model: nn.Module) -> float:\n",
    "    model.eval()  # turn on evaluation mode\n",
    "    total_loss = 0.\n",
    "    # src_mask = generate_square_subsequent_mask(bptt).to(device)\n",
    "    with torch.no_grad():\n",
    "        predictions = []\n",
    "        true_labels = []\n",
    "        for batch, i in enumerate(val_dataloader):\n",
    "            data, mask, targets = i.values()\n",
    "            data = data.to(device)\n",
    "            mask = mask.to(device)\n",
    "            targets = targets.to(device)\n",
    "            seq_len = data.size(1)\n",
    "            # logits_task1, logits_task2, logits_task3, ae_output = model(data, mask)\n",
    "            # task_logits, ae_output = model(data)\n",
    "            # task_logits, ae_output = model(data, mask)\n",
    "            output = model(data, mask)\n",
    "            # t1_out, t2_out, t3_out, auto_output = model(data, t1, t2, t3)\n",
    "            # loss = custom_loss(logits_task1, logits_task2, logits_task3, targets)\n",
    "            # loss = custom_loss(logits_task1, logits_task2, logits_task3, ae_output, data, targets)\n",
    "            loss = criterion(output, targets.flatten())\n",
    "\n",
    "            total_loss += seq_len * loss.item()\n",
    "\n",
    "            #get the labels for classification report\n",
    "            pred_label, true_label = get_labels(output, targets)\n",
    "            predictions.extend(pred_label)\n",
    "            true_labels.extend(true_label)\n",
    "            # if batch == 100:\n",
    "            #     break\n",
    "\n",
    "    # Compute overall classification report\n",
    "    logging.info(f\"\\n Scores:\")\n",
    "    acc = accuracy_score(true_labels, predictions)\n",
    "    logging.info(f\"\\n Accuracy:{acc}\")\n",
    "    logging.info(f\"\\n {classification_report(true_labels, predictions)}\")\n",
    "    return total_loss / (len(val_dataset) - 1), acc\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-07-13 12:37:38,842 - INFO - #########################################################################################\n",
      "2023-07-13 12:37:38,844 - INFO - \n",
      " DESCRIPTION-> \n",
      " logic: XLMR(finetuning last 3-layers) + linear_layer + loss_reweighting (epochs=20), model: xlm-roberta-base, lr:0.0002, max_seq_length: 500\n",
      "2023-07-13 12:37:38,845 - INFO - #########################################################################################\n"
     ]
    }
   ],
   "source": [
    "logging.info(f\"#\"* 89)\n",
    "logging.info(f\"\\n DESCRIPTION-> \\n logic: XLMR(finetuning last 3-layers) + linear_layer + loss_reweighting (epochs=20), model: {tokenizer.name_or_path}, lr:{learning_rate}, max_seq_length: {max_length}\")\n",
    "logging.info('#' * 89)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▎         | 100/2732 [03:37<1:35:22,  2.17s/it]\n",
      "100%|██████████| 182/182 [07:48<00:00,  2.57s/it]\n",
      "2023-07-13 12:45:22,701 - INFO - \n",
      " Scores:\n",
      "2023-07-13 12:45:22,713 - INFO - \n",
      " Accuracy:0.19414182088975923\n",
      "/data1/debajyoti/colie/colenv/lib/python3.8/site-packages/sklearn/metrics/_classification.py:1469: 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",
      "/data1/debajyoti/colie/colenv/lib/python3.8/site-packages/sklearn/metrics/_classification.py:1469: 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",
      "/data1/debajyoti/colie/colenv/lib/python3.8/site-packages/sklearn/metrics/_classification.py:1469: 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",
      "2023-07-13 12:45:22,771 - INFO - \n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "           0       0.10      0.74      0.18      1158\n",
      "           1       0.73      0.05      0.09     16938\n",
      "           2       0.69      0.29      0.41     14848\n",
      "           3       0.00      0.00      0.00      1713\n",
      "           4       0.05      0.70      0.10      1600\n",
      "\n",
      "    accuracy                           0.19     36257\n",
      "   macro avg       0.32      0.35      0.16     36257\n",
      "weighted avg       0.63      0.19      0.22     36257\n",
      "\n",
      "2023-07-13 12:45:22,772 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 12:45:22,773 - INFO - | end of epoch   1 | time: 463.92s | valid loss 3.73663\n",
      "2023-07-13 12:45:22,774 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 12:53:04,945 - INFO - \n",
      " Scores:\n",
      "2023-07-13 12:53:04,963 - INFO - \n",
      " Accuracy:0.3259232699892435\n",
      "2023-07-13 12:53:05,010 - INFO - \n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "           0       0.13      0.58      0.22      1158\n",
      "           1       0.74      0.24      0.36     16938\n",
      "           2       0.55      0.42      0.47     14848\n",
      "           3       0.27      0.21      0.24      1713\n",
      "           4       0.04      0.35      0.07      1600\n",
      "\n",
      "    accuracy                           0.33     36257\n",
      "   macro avg       0.35      0.36      0.27     36257\n",
      "weighted avg       0.59      0.33      0.38     36257\n",
      "\n",
      "2023-07-13 12:53:05,011 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 12:53:05,012 - INFO - | end of epoch   2 | time: 460.74s | valid loss 3.79886\n",
      "2023-07-13 12:53:05,013 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:00:58,594 - INFO - \n",
      " Scores:\n",
      "2023-07-13 13:00:58,609 - INFO - \n",
      " Accuracy:0.27994594147337065\n",
      "2023-07-13 13:00:58,650 - INFO - \n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "           0       0.12      0.51      0.20      1158\n",
      "           1       0.79      0.08      0.15     16938\n",
      "           2       0.68      0.47      0.56     14848\n",
      "           3       0.30      0.21      0.25      1713\n",
      "           4       0.04      0.49      0.08      1600\n",
      "\n",
      "    accuracy                           0.28     36257\n",
      "   macro avg       0.39      0.35      0.25     36257\n",
      "weighted avg       0.67      0.28      0.32     36257\n",
      "\n",
      "2023-07-13 13:00:58,651 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:00:58,652 - INFO - | end of epoch   3 | time: 462.55s | valid loss 3.64016\n",
      "2023-07-13 13:00:58,653 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:08:39,738 - INFO - \n",
      " Scores:\n",
      "2023-07-13 13:08:39,756 - INFO - \n",
      " Accuracy:0.382767465592851\n",
      "2023-07-13 13:08:39,800 - INFO - \n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "           0       0.12      0.61      0.21      1158\n",
      "           1       0.71      0.30      0.42     16938\n",
      "           2       0.71      0.47      0.57     14848\n",
      "           3       0.30      0.21      0.24      1713\n",
      "           4       0.06      0.44      0.10      1600\n",
      "\n",
      "    accuracy                           0.38     36257\n",
      "   macro avg       0.38      0.41      0.31     36257\n",
      "weighted avg       0.64      0.38      0.45     36257\n",
      "\n",
      "2023-07-13 13:08:39,802 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:08:39,804 - INFO - | end of epoch   4 | time: 461.15s | valid loss 3.61403\n",
      "2023-07-13 13:08:39,805 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:16:34,748 - INFO - \n",
      " Scores:\n",
      "2023-07-13 13:16:34,767 - INFO - \n",
      " Accuracy:0.4561050279945941\n",
      "2023-07-13 13:16:34,806 - INFO - \n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "           0       0.17      0.26      0.20      1158\n",
      "           1       0.71      0.53      0.61     16938\n",
      "           2       0.71      0.42      0.53     14848\n",
      "           3       0.25      0.27      0.26      1713\n",
      "           4       0.04      0.29      0.07      1600\n",
      "\n",
      "    accuracy                           0.46     36257\n",
      "   macro avg       0.37      0.36      0.33     36257\n",
      "weighted avg       0.64      0.46      0.52     36257\n",
      "\n",
      "2023-07-13 13:16:34,808 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:16:34,809 - INFO - | end of epoch   5 | time: 462.29s | valid loss 3.35628\n",
      "2023-07-13 13:16:34,810 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:24:29,198 - INFO - \n",
      " Scores:\n",
      "2023-07-13 13:24:29,211 - INFO - \n",
      " Accuracy:0.2591775381305679\n",
      "2023-07-13 13:24:29,270 - INFO - \n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "           0       0.13      0.63      0.21      1158\n",
      "           1       0.76      0.16      0.26     16938\n",
      "           2       0.67      0.32      0.44     14848\n",
      "           3       0.28      0.29      0.29      1713\n",
      "           4       0.04      0.42      0.07      1600\n",
      "\n",
      "    accuracy                           0.26     36257\n",
      "   macro avg       0.38      0.36      0.25     36257\n",
      "weighted avg       0.65      0.26      0.33     36257\n",
      "\n",
      "2023-07-13 13:24:29,272 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:24:29,273 - INFO - | end of epoch   6 | time: 463.23s | valid loss 3.80452\n",
      "2023-07-13 13:24:29,274 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:32:12,135 - INFO - \n",
      " Scores:\n",
      "2023-07-13 13:32:12,148 - INFO - \n",
      " Accuracy:0.44669994759632625\n",
      "2023-07-13 13:32:12,194 - INFO - \n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "           0       0.16      0.53      0.24      1158\n",
      "           1       0.75      0.40      0.52     16938\n",
      "           2       0.70      0.53      0.60     14848\n",
      "           3       0.32      0.26      0.29      1713\n",
      "           4       0.05      0.32      0.08      1600\n",
      "\n",
      "    accuracy                           0.45     36257\n",
      "   macro avg       0.39      0.41      0.35     36257\n",
      "weighted avg       0.66      0.45      0.52     36257\n",
      "\n",
      "2023-07-13 13:32:12,195 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:32:12,196 - INFO - | end of epoch   7 | time: 462.92s | valid loss 3.33110\n",
      "2023-07-13 13:32:12,197 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:39:56,084 - INFO - \n",
      " Scores:\n",
      "2023-07-13 13:39:56,097 - INFO - \n",
      " Accuracy:0.3827398847119177\n",
      "2023-07-13 13:39:56,136 - INFO - \n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "           0       0.15      0.56      0.24      1158\n",
      "           1       0.76      0.32      0.45     16938\n",
      "           2       0.72      0.44      0.55     14848\n",
      "           3       0.28      0.35      0.31      1713\n",
      "           4       0.04      0.37      0.08      1600\n",
      "\n",
      "    accuracy                           0.38     36257\n",
      "   macro avg       0.39      0.41      0.33     36257\n",
      "weighted avg       0.67      0.38      0.46     36257\n",
      "\n",
      "2023-07-13 13:39:56,138 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:39:56,139 - INFO - | end of epoch   8 | time: 463.94s | valid loss 3.51282\n",
      "2023-07-13 13:39:56,139 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:47:38,915 - INFO - \n",
      " Scores:\n",
      "2023-07-13 13:47:38,929 - INFO - \n",
      " Accuracy:0.42300797087458974\n",
      "2023-07-13 13:47:38,974 - INFO - \n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "           0       0.12      0.66      0.21      1158\n",
      "           1       0.71      0.45      0.56     16938\n",
      "           2       0.78      0.39      0.52     14848\n",
      "           3       0.35      0.29      0.31      1713\n",
      "           4       0.06      0.40      0.11      1600\n",
      "\n",
      "    accuracy                           0.42     36257\n",
      "   macro avg       0.40      0.44      0.34     36257\n",
      "weighted avg       0.68      0.42      0.50     36257\n",
      "\n",
      "2023-07-13 13:47:38,976 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:47:38,977 - INFO - | end of epoch   9 | time: 462.84s | valid loss 3.42163\n",
      "2023-07-13 13:47:38,979 - INFO - -----------------------------------------------------------------------------------------\n",
      "2023-07-13 13:47:38,980 - INFO - Early stopped training at epoch 9\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_val_acc = 0\n",
    "current_val_loss = []   # for plotting graph of val_loss\n",
    "epochs = 50\n",
    "early_stop_thresh = 3\n",
    "\n",
    "tempdir = '/data1/debajyoti/colie/.temp/'\n",
    "best_model_params_path = os.path.join(tempdir, f\"best_model_params_{time.asctime().replace(' ','_')}.pt\")\n",
    "\n",
    "for epoch in range(1, epochs + 1):\n",
    "    epoch_start_time = time.time()\n",
    "    train(model)\n",
    "    val_loss, accuracy = evaluate(model)\n",
    "    current_val_loss.append(val_loss)\n",
    "    # val_ppl = np.exp(val_loss)\n",
    "    elapsed = time.time() - epoch_start_time\n",
    "    logging.info('-' * 89)\n",
    "    logging.info(f'| end of epoch {epoch:3d} | time: {elapsed:5.2f}s | '\n",
    "        f'valid loss {val_loss:5.5f}')\n",
    "    logging.info('-' * 89)\n",
    "\n",
    "    if accuracy > best_val_acc:\n",
    "        best_val_acc = accuracy\n",
    "        best_epoch = epoch\n",
    "        torch.save(model.state_dict(), best_model_params_path)\n",
    "    elif epoch - best_epoch > early_stop_thresh:\n",
    "        logging.info(\"Early stopped training at epoch %d\" % epoch)\n",
    "        break  # terminate the training loop\n",
    "\n",
    "    scheduler.step()\n",
    "model.load_state_dict(torch.load(best_model_params_path)) # load best model states\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "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>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>\"Alas, poor girl!\" said I, \"I fear that her ha...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>to divide her attention between the said garco...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>visitor's disposition to gallantry. However, s...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>says Juvenal, \"'Mors sola fatetur Quantula sin...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>him out in that back passage; the outer door i...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143025</th>\n",
       "      <td>be hard for anyone to do anything dignified, w...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143026</th>\n",
       "      <td>Wilson, and could not bring himself to think t...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143027</th>\n",
       "      <td>give them a chance, the same as everybody else...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143028</th>\n",
       "      <td>political convention in which he declared that...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143029</th>\n",
       "      <td>Jimmie doesn't know anything about the Russian...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>143030 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                     text\n",
       "0       \"Alas, poor girl!\" said I, \"I fear that her ha...\n",
       "1       to divide her attention between the said garco...\n",
       "2       visitor's disposition to gallantry. However, s...\n",
       "3       says Juvenal, \"'Mors sola fatetur Quantula sin...\n",
       "4       him out in that back passage; the outer door i...\n",
       "...                                                   ...\n",
       "143025  be hard for anyone to do anything dignified, w...\n",
       "143026  Wilson, and could not bring himself to think t...\n",
       "143027  give them a chance, the same as everybody else...\n",
       "143028  political convention in which he declared that...\n",
       "143029  Jimmie doesn't know anything about the Russian...\n",
       "\n",
       "[143030 rows x 1 columns]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/716 [00:00<?, ?it/s]"
     ]
    }
   ],
   "source": [
    "class CustomDataset_test(Dataset):\n",
    "    def __init__(self, tokenizer, df):\n",
    "        # Initialize the tokenizer\n",
    "        self.tokenizer = tokenizer\n",
    "\n",
    "        self.df = df\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        # Get the text and label from the dataframe\n",
    "        text = self.df.iloc[index]['text']\n",
    "        # label = self.df.iloc[index]['label']\n",
    "\n",
    "        # Tokenize the text and convert it to input IDs\n",
    "        inputs = self.tokenizer(\n",
    "            text,\n",
    "            None,\n",
    "            add_special_tokens=False,\n",
    "            max_length=max_length,\n",
    "            padding='max_length',\n",
    "            truncation=True,\n",
    "            return_token_type_ids=False,\n",
    "            return_attention_mask=True,\n",
    "            return_tensors='pt',\n",
    "        )\n",
    "\n",
    "\n",
    "        # Return the input IDs and label as PyTorch tensors\n",
    "        return {\n",
    "            'input_ids': inputs['input_ids'][0],\n",
    "            'attention_mask': inputs['attention_mask'][0],\n",
    "            # 'token_type_ids': inputs['token_type_ids'][0],\n",
    "            # 'label': torch.tensor(label, dtype=torch.int64),\n",
    "        }\n",
    "\n",
    "test_dataset = CustomDataset_test(tokenizer, test)\n",
    "\n",
    "# DataLoader\n",
    "test_dataloader = tqdm(DataLoader(test_dataset, batch_size=batch_size, shuffle=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 716/716 [26:22<00:00,  2.21s/it]\n"
     ]
    }
   ],
   "source": [
    "predictions = []\n",
    "\n",
    "# Evaluate model on test data\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for batch_idx, batch in enumerate(test_dataloader):\n",
    "        data, mask = batch.values()\n",
    "        data = data.to(device)\n",
    "        mask = mask.to(device)\n",
    "        # targets = targets.to(device)\n",
    "        seq_len = data.size(1)\n",
    "        # logits_task1, logits_task2, logits_task3, ae_output = model(data, mask)\n",
    "        # task_logits, ae_output = model(data)\n",
    "        # task_logits, ae_output = model(data, mask)\n",
    "        output = model(data, mask)\n",
    "        #get the labels for classification report\n",
    "        pred_label = torch.argmax(output, dim=1).flatten().tolist()\n",
    "        predictions.extend(pred_label)\n",
    "        \n",
    "        # if batch == 400:\n",
    "        #     break\n",
    "\n",
    "# # Compute overall classification report\n",
    "# logging.info(f\"\\n Scores:\")\n",
    "# logging.info(f\"\\n {classification_report(true_labels, predictions)}\")\n",
    "# return total_loss / (len(val_dataset) - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_labels[\"Epoch\"] = predictions\n",
    "# test_labels.to_csv('file_name.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "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>BOOK_id</th>\n",
       "      <th>Epoch</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>7616_1.txt</td>\n",
       "      <td>Romanticism</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>7616_2.txt</td>\n",
       "      <td>Viktorian</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>7616_3.txt</td>\n",
       "      <td>Viktorian</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7616_4.txt</td>\n",
       "      <td>Viktorian</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>7616_5.txt</td>\n",
       "      <td>Romanticism</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143025</th>\n",
       "      <td>5677_92.txt</td>\n",
       "      <td>OurDays</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143026</th>\n",
       "      <td>5677_93.txt</td>\n",
       "      <td>OurDays</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143027</th>\n",
       "      <td>5677_94.txt</td>\n",
       "      <td>OurDays</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143028</th>\n",
       "      <td>5677_95.txt</td>\n",
       "      <td>OurDays</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>143029</th>\n",
       "      <td>5677_96.txt</td>\n",
       "      <td>OurDays</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>143030 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            BOOK_id        Epoch\n",
       "0        7616_1.txt  Romanticism\n",
       "1        7616_2.txt    Viktorian\n",
       "2        7616_3.txt    Viktorian\n",
       "3        7616_4.txt    Viktorian\n",
       "4        7616_5.txt  Romanticism\n",
       "...             ...          ...\n",
       "143025  5677_92.txt      OurDays\n",
       "143026  5677_93.txt      OurDays\n",
       "143027  5677_94.txt      OurDays\n",
       "143028  5677_95.txt      OurDays\n",
       "143029  5677_96.txt      OurDays\n",
       "\n",
       "[143030 rows x 2 columns]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "label_dic = {0:'Romanticism',\n",
    "            1:'Viktorian',\n",
    "            2:'Modernism',\n",
    "            3:'PostModernism',\n",
    "            4:'OurDays'}\n",
    "test_labels['Epoch'] = test_labels['Epoch'].map(label_dic)\n",
    "test_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_labels.to_csv('submission_2.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}