multitask_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "120740cc-1ba9-44ae-a673-1047f23de335",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import transformers\n",
    "transformers.logging.set_verbosity_error()\n",
    "from transformers import RobertaTokenizer, RobertaForSequenceClassification, Trainer, TrainingArguments\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, DataCollatorWithPadding\n",
    "from datasets import load_dataset, DatasetDict, Dataset, load_metric\n",
    "from sklearn.metrics import confusion_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "18cf2097",
   "metadata": {},
   "outputs": [],
   "source": [
    "from classes import RobertaForSequenceClassification2, StrIgnoreDevice, DataLoaderWithTaskname, MultitaskDataloader, MultitaskTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "17b7a002",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = RobertaTokenizer.from_pretrained(\"roberta-base\")\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"text\"], truncation=True)\n",
    "\n",
    "def read_and_tokenize_data(train,test,evaluate):\n",
    "    d = load_dataset('pandas', data_files={\"train\":train, \"test\":test,\"evaluate\":evaluate})\n",
    "    for i in d:\n",
    "        d[i]= d[i].remove_columns(column_names = ['__index_level_0__'])\n",
    "    d = d.map(tokenize_function, batched=True)\n",
    "    return d\n",
    "\n",
    "data_collator = DataCollatorWithPadding(tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "03489c5b",
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration default-5d8442b82fe9acdc\n",
      "Reusing dataset pandas (C:\\Users\\dovyd\\.cache\\huggingface\\datasets\\pandas\\default-5d8442b82fe9acdc\\0.0.0\\6197c1e855b639d75a767140856841a562b7a71d129104973fe1962594877ade)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f5a9b4bfb3f442ada77aa7418e32dc5d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at C:\\Users\\dovyd\\.cache\\huggingface\\datasets\\pandas\\default-5d8442b82fe9acdc\\0.0.0\\6197c1e855b639d75a767140856841a562b7a71d129104973fe1962594877ade\\cache-c986a07d4e839639.arrow\n",
      "Loading cached processed dataset at C:\\Users\\dovyd\\.cache\\huggingface\\datasets\\pandas\\default-5d8442b82fe9acdc\\0.0.0\\6197c1e855b639d75a767140856841a562b7a71d129104973fe1962594877ade\\cache-1aef789b7c6c8461.arrow\n",
      "Loading cached processed dataset at C:\\Users\\dovyd\\.cache\\huggingface\\datasets\\pandas\\default-5d8442b82fe9acdc\\0.0.0\\6197c1e855b639d75a767140856841a562b7a71d129104973fe1962594877ade\\cache-4ecbd960b02fdce5.arrow\n",
      "Using custom data configuration default-d0f7ab1575e159e4\n",
      "Reusing dataset pandas (C:\\Users\\dovyd\\.cache\\huggingface\\datasets\\pandas\\default-d0f7ab1575e159e4\\0.0.0\\6197c1e855b639d75a767140856841a562b7a71d129104973fe1962594877ade)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "22a203afb57c4313b762e27d1a74e882",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at C:\\Users\\dovyd\\.cache\\huggingface\\datasets\\pandas\\default-d0f7ab1575e159e4\\0.0.0\\6197c1e855b639d75a767140856841a562b7a71d129104973fe1962594877ade\\cache-1ddf9eac46f9543e.arrow\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2e6cd2cb4cea4c54a3ed571651104c0e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c0d0882c642460883b930e85f3d8add",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset_dict = {\n",
    "    \"author\":  read_and_tokenize_data(\"data/authors_train.pkl\",\"data/authors_test.pkl\",\"data/authors_validation.pkl\"),\n",
    "    \"sentiment\":  read_and_tokenize_data(\"data/imdb_train.pkl\",\"data/imdb_test.pkl\",\"data/imdb_validation.pkl\")\n",
    "}\n",
    "\n",
    "\n",
    "multitask_model = RobertaForSequenceClassification2.from_pretrained(\n",
    "        \"roberta-base\",\n",
    "        task_labels_map={\"author\": 4, \"sentiment\": 2},\n",
    "    )\n",
    "\n",
    "\n",
    "for i in dataset_dict.values():\n",
    "    for j in i.values():\n",
    "        j.set_format(type=\"torch\",columns=[\"input_ids\", \"attention_mask\", \"label\"])\n",
    "    \n",
    "    \n",
    "train_dataset = {\n",
    "        task_name: dataset[\"train\"] for task_name, dataset in dataset_dict.items()}\n",
    "\n",
    "eval_dataset = {\n",
    "        task_name: dataset[\"evaluate\"] for task_name, dataset in dataset_dict.items()}\n",
    "\n",
    "test_dataset = {\n",
    "        task_name: dataset[\"test\"] for task_name, dataset in dataset_dict.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b346fc3f",
   "metadata": {},
   "source": [
    "#### Training the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1c6c7d7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir='results',\n",
    "    learning_rate=2e-5,\n",
    "    weight_decay=0.01,\n",
    "    num_train_epochs=15,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps = 500,\n",
    "    save_steps = 500,\n",
    "    load_best_model_at_end = True,\n",
    "    per_device_eval_batch_size =  8\n",
    ")\n",
    "\n",
    "\n",
    "metric = load_metric(\"accuracy\")\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    logits, labels = eval_pred\n",
    "    predictions = np.argmax(logits, axis=-1)\n",
    "    return metric.compute(predictions=predictions, references=labels)\n",
    "\n",
    "\n",
    "data_collator = DataCollatorWithPadding(tokenizer=tokenizer)\n",
    "\n",
    "trainer = MultitaskTrainer(\n",
    "    model=multitask_model,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=eval_dataset,\n",
    "    args=training_args,\n",
    "    compute_metrics=compute_metrics,\n",
    "    data_collator=data_collator\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ad4b49f7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "***** Running training *****\n",
      "  Num examples = 5391\n",
      "  Num Epochs = 15\n",
      "  Instantaneous batch size per device = 8\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 8\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 10125\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6787, 'learning_rate': 1.901234567901235e-05, 'epoch': 0.74}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-500\n",
      "Configuration saved in results/checkpoint-500/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5599084496498108, 'eval_accuracy': 0.8466666666666667, 'eval_runtime': 3.2378, 'eval_samples_per_second': 92.656, 'eval_steps_per_second': 11.736, 'epoch': 0.74}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-500/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4569, 'learning_rate': 1.802469135802469e-05, 'epoch': 1.48}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-1000\n",
      "Configuration saved in results/checkpoint-1000/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.48681876063346863, 'eval_accuracy': 0.8766666666666667, 'eval_runtime': 3.0574, 'eval_samples_per_second': 98.124, 'eval_steps_per_second': 12.429, 'epoch': 1.48}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-1000/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.3162, 'learning_rate': 1.7037037037037038e-05, 'epoch': 2.22}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-1500\n",
      "Configuration saved in results/checkpoint-1500/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.37365713715553284, 'eval_accuracy': 0.9033333333333333, 'eval_runtime': 3.0387, 'eval_samples_per_second': 98.728, 'eval_steps_per_second': 12.506, 'epoch': 2.22}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-1500/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.1991, 'learning_rate': 1.6049382716049385e-05, 'epoch': 2.96}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-2000\n",
      "Configuration saved in results/checkpoint-2000/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.4169445037841797, 'eval_accuracy': 0.9333333333333333, 'eval_runtime': 3.0751, 'eval_samples_per_second': 97.558, 'eval_steps_per_second': 12.357, 'epoch': 2.96}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-2000/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.1429, 'learning_rate': 1.506172839506173e-05, 'epoch': 3.7}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-2500\n",
      "Configuration saved in results/checkpoint-2500/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.42519184947013855, 'eval_accuracy': 0.92, 'eval_runtime': 3.0626, 'eval_samples_per_second': 97.955, 'eval_steps_per_second': 12.408, 'epoch': 3.7}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-2500/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0857, 'learning_rate': 1.4074074074074075e-05, 'epoch': 4.44}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-3000\n",
      "Configuration saved in results/checkpoint-3000/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5788098573684692, 'eval_accuracy': 0.9166666666666666, 'eval_runtime': 3.1225, 'eval_samples_per_second': 96.076, 'eval_steps_per_second': 12.17, 'epoch': 4.44}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-3000/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0634, 'learning_rate': 1.3086419753086422e-05, 'epoch': 5.19}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-3500\n",
      "Configuration saved in results/checkpoint-3500/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.4006443917751312, 'eval_accuracy': 0.94, 'eval_runtime': 3.05, 'eval_samples_per_second': 98.362, 'eval_steps_per_second': 12.459, 'epoch': 5.19}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-3500/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0424, 'learning_rate': 1.2098765432098767e-05, 'epoch': 5.93}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-4000\n",
      "Configuration saved in results/checkpoint-4000/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5864290595054626, 'eval_accuracy': 0.93, 'eval_runtime': 3.0659, 'eval_samples_per_second': 97.852, 'eval_steps_per_second': 12.395, 'epoch': 5.93}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-4000/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0423, 'learning_rate': 1.1111111111111113e-05, 'epoch': 6.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-4500\n",
      "Configuration saved in results/checkpoint-4500/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.41267475485801697, 'eval_accuracy': 0.9466666666666667, 'eval_runtime': 3.0452, 'eval_samples_per_second': 98.517, 'eval_steps_per_second': 12.479, 'epoch': 6.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-4500/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0304, 'learning_rate': 1.0123456790123458e-05, 'epoch': 7.41}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-5000\n",
      "Configuration saved in results/checkpoint-5000/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.6471032500267029, 'eval_accuracy': 0.92, 'eval_runtime': 3.0773, 'eval_samples_per_second': 97.49, 'eval_steps_per_second': 12.349, 'epoch': 7.41}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-5000/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0129, 'learning_rate': 9.135802469135803e-06, 'epoch': 8.15}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-5500\n",
      "Configuration saved in results/checkpoint-5500/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.6132160425186157, 'eval_accuracy': 0.9266666666666666, 'eval_runtime': 3.0603, 'eval_samples_per_second': 98.028, 'eval_steps_per_second': 12.417, 'epoch': 8.15}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-5500/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0311, 'learning_rate': 8.148148148148148e-06, 'epoch': 8.89}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-6000\n",
      "Configuration saved in results/checkpoint-6000/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5779215693473816, 'eval_accuracy': 0.9333333333333333, 'eval_runtime': 3.0676, 'eval_samples_per_second': 97.797, 'eval_steps_per_second': 12.388, 'epoch': 8.89}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-6000/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0133, 'learning_rate': 7.160493827160494e-06, 'epoch': 9.63}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-6500\n",
      "Configuration saved in results/checkpoint-6500/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5070252418518066, 'eval_accuracy': 0.94, 'eval_runtime': 3.0835, 'eval_samples_per_second': 97.291, 'eval_steps_per_second': 12.324, 'epoch': 9.63}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-6500/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0299, 'learning_rate': 6.17283950617284e-06, 'epoch': 10.37}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-7000\n",
      "Configuration saved in results/checkpoint-7000/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.6678612232208252, 'eval_accuracy': 0.9266666666666666, 'eval_runtime': 3.103, 'eval_samples_per_second': 96.681, 'eval_steps_per_second': 12.246, 'epoch': 10.37}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-7000/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0109, 'learning_rate': 5.185185185185185e-06, 'epoch': 11.11}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-7500\n",
      "Configuration saved in results/checkpoint-7500/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5488567352294922, 'eval_accuracy': 0.9366666666666666, 'eval_runtime': 3.105, 'eval_samples_per_second': 96.619, 'eval_steps_per_second': 12.238, 'epoch': 11.11}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-7500/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0018, 'learning_rate': 4.197530864197531e-06, 'epoch': 11.85}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-8000\n",
      "Configuration saved in results/checkpoint-8000/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.6135303974151611, 'eval_accuracy': 0.9333333333333333, 'eval_runtime': 3.0652, 'eval_samples_per_second': 97.872, 'eval_steps_per_second': 12.397, 'epoch': 11.85}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-8000/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0114, 'learning_rate': 3.2098765432098767e-06, 'epoch': 12.59}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-8500\n",
      "Configuration saved in results/checkpoint-8500/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.578200101852417, 'eval_accuracy': 0.9433333333333334, 'eval_runtime': 3.1536, 'eval_samples_per_second': 95.129, 'eval_steps_per_second': 12.05, 'epoch': 12.59}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-8500/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0, 'learning_rate': 2.222222222222222e-06, 'epoch': 13.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-9000\n",
      "Configuration saved in results/checkpoint-9000/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5728410482406616, 'eval_accuracy': 0.9433333333333334, 'eval_runtime': 3.0618, 'eval_samples_per_second': 97.98, 'eval_steps_per_second': 12.411, 'epoch': 13.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-9000/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0, 'learning_rate': 1.234567901234568e-06, 'epoch': 14.07}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-9500\n",
      "Configuration saved in results/checkpoint-9500/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5762560367584229, 'eval_accuracy': 0.9433333333333334, 'eval_runtime': 3.0501, 'eval_samples_per_second': 98.358, 'eval_steps_per_second': 12.459, 'epoch': 14.07}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-9500/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 300\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.0, 'learning_rate': 2.469135802469136e-07, 'epoch': 14.81}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to results/checkpoint-10000\n",
      "Configuration saved in results/checkpoint-10000/config.json\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.6036664843559265, 'eval_accuracy': 0.94, 'eval_runtime': 3.0673, 'eval_samples_per_second': 97.806, 'eval_steps_per_second': 12.389, 'epoch': 14.81}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model weights saved in results/checkpoint-10000/pytorch_model.bin\n",
      "\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n",
      "Loading best model from results/checkpoint-1500 (score: 0.37365713715553284).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'train_runtime': 2821.0705, 'train_samples_per_second': 0.011, 'train_steps_per_second': 3.589, 'train_loss': 0.10714140081626397, 'epoch': 15.0}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=10125, training_loss=0.10714140081626397, metrics={'train_runtime': 2821.0705, 'train_samples_per_second': 0.011, 'train_steps_per_second': 3.589, 'train_loss': 0.10714140081626397, 'epoch': 15.0})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d484c224",
   "metadata": {},
   "source": [
    "#### Saving the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fd2cb060",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to models/multihead_classification/model\n",
      "Configuration saved in models/multihead_classification/model/config.json\n",
      "Model weights saved in models/multihead_classification/model/pytorch_model.bin\n",
      "tokenizer config file saved in models/multihead_classification/tokenizer/tokenizer_config.json\n",
      "Special tokens file saved in models/multihead_classification/tokenizer/special_tokens_map.json\n"
     ]
    }
   ],
   "source": [
    "def save_model(trainer,tokenizer,trainer_path,tokenizer_path):\n",
    "    trainer.save_model(trainer_path)\n",
    "    tokenizer.save_pretrained(tokenizer_path)\n",
    "    \n",
    "save_model(trainer,tokenizer,\"models/multihead_classification/model\",\"models/multihead_classification/tokenizer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "daac0799",
   "metadata": {},
   "source": [
    "#### Loading the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d3f66af0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model(model_path,tokenizer_path):\n",
    "    model = RobertaForSequenceClassification2.from_pretrained(model_path,task_labels_map={\"author\": 4, \"sentiment\": 2})\n",
    "    tokenizer = RobertaTokenizer.from_pretrained(tokenizer_path)\n",
    "    trainer = Trainer(model = model)\n",
    "    return (model, trainer, tokenizer)\n",
    "    \n",
    "\n",
    "model, trainer, tokenizer = load_model(\"models/multihead_classification/model\",\"models/multihead_classification/tokenizer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2f13a9c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "def make_predictions(test_data):\n",
    "    metric = load_metric(\"accuracy\")\n",
    "    for task in [\"author\", \"sentiment\"]:\n",
    "            y_true = []\n",
    "            y_pred = []\n",
    "            for i in test_data[task]:\n",
    "                logits = model(input_ids = torch.unsqueeze(i[\"input_ids\"],0).to(device),\\\n",
    "                                         attention_mask = torch.unsqueeze(i[\"attention_mask\"],0).to(device),task_name=task)[0]\n",
    "                y_true.append(i[\"label\"])\n",
    "                y_pred.append(np.argmax(logits.cpu().detach().numpy(), axis=-1))\n",
    "            print(confusion_matrix(y_true, y_pred))\n",
    "            acc = metric.compute(predictions=y_pred, references=y_true)\n",
    "            print(f\"Task: {task}\", f\"Accuracy: {acc}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "403bc7e2-7033-46e6-a9ab-351776f45a5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[48  0  0  0]\n",
      " [ 1 20  0  0]\n",
      " [ 1  0  4  1]\n",
      " [ 2  8  0 64]]\n",
      "Task: author Accuracy: {'accuracy': 0.912751677852349}\n",
      "[[70  7]\n",
      " [ 4 69]]\n",
      "Task: sentiment Accuracy: {'accuracy': 0.9266666666666666}\n"
     ]
    }
   ],
   "source": [
    "make_predictions(test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ae27cf7b-a816-42a0-8a21-f44fdac79aeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pipeline(text, task):\n",
    "    inputs = tokenizer(text,padding=\"longest\",return_tensors=\"pt\")\n",
    "    inputs[\"input_ids\"] =  inputs[\"input_ids\"].to(device)\n",
    "    inputs[\"attention_mask\"] =  inputs[\"attention_mask\"].to(device)\n",
    "    logits = model(**inputs, task_name=task)[\"logits\"]\n",
    "    probabilities = torch.softmax(logits, dim=1).detach().cpu()\n",
    "    result = torch.argmax(\n",
    "        probabilities,axis=1\n",
    "    )\n",
    "    probability = torch.max(\n",
    "        probabilities,axis=1\n",
    "    )[0]\n",
    "    print({\"label\":result, \"score\":probability})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0fd2cd1d-1ce0-4ae4-9aa9-47e4f06ae760",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'label': tensor([0]), 'score': tensor([0.9992])}\n"
     ]
    }
   ],
   "source": [
    "pipeline(\"Total dissapointment. I do not understand why would anyone watch this\",\"sentiment\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "02c5fea5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'label': tensor([1]), 'score': tensor([0.9981])}\n"
     ]
    }
   ],
   "source": [
    "pipeline(\"An all time classic\",\"sentiment\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}