bert_multitask_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchaudio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#check for GPU\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pickle\n",
    "import json\n",
    "import random\n",
    "import time\n",
    "from transformers import BertTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_relations = {'Comment':0, 'Contrast':1, 'Correction':2, 'Question-answer_pair':3, 'Acknowledgement':4,'Elaboration':5,\n",
    "                 'Clarification_question':6, 'Conditional':7, 'Continuation':8, 'Result':9, 'Explanation':10, 'Q-Elab':11,\n",
    "                 'Alternation':12, 'Narration':13, 'Confirmation_question':14, 'Sequence':15, 'Break':16}# "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reverse_relations = {0:'Comment', 1:'Contrast', 2:'Correction', 3:'Question-answer_pair', 4:'Acknowledgement',5:'Elaboration',\n",
    "                 6:'Clarification_question', 7:'Conditional', 8:'Continuation', 9:'Result', 10:'Explanation', 11:'Q-Elab',\n",
    "                 12:'Alternation', 13:'Narration', 14:'Confirmation_question', 15:'Sequence', 16:'Break'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "home=%pwd\n",
    "filename = home + '/data/TEST_101_bert.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import load_data, input_format, position_ids_compute, tokenize\n",
    "from bert_format import undersample, format_time, flat_accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = load_data(filename, map_relations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained('bert-base-cased', use_fast=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "put = ['1','0']\n",
    "colors = ['r', 'b', 'g', 'o', 'y', 'p']\n",
    "listx = ['b', 'c', 'd', 'f', 'g', 'h', 'j', 'k', 'l', 'm', 'n']\n",
    "listy = ['0', '1', '2', '3', '4', '5', '6', '7', '8']\n",
    "listz = ['a', 'e', 'i', 'o', 'u', 'p', 'q', 'r', 'x', 'y', 'z']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coord_tokens = [''.join([s, t, i, j, k]) for s in put\n",
    "                for t in colors\n",
    "                for i in listx\n",
    "                for j in listy\n",
    "                for k in listz]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.add_tokens(coord_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs, labels_input, raw = input_format(test_data, 10, relations=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels_input[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_labels = len(set([r[3] for r in labels_input]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_labels = 17"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_tokenized = tokenizer(inputs, return_tensors=\"pt\", padding=True, truncation=True, add_special_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_ids = batch_tokenized[\"input_ids\"].to(device) # list of token ids of dialogs in batch\n",
    "attention_masks = batch_tokenized[\"attention_mask\"].to(device)\n",
    "token_type_ids = batch_tokenized[\"token_type_ids\"].to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = [label[3] for label in list(labels_input)]\n",
    "labels = torch.tensor(labels)\n",
    "labels_relation = torch.tensor(labels_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "position_ids = position_ids_compute(tokenizer, input_ids, raw, labels_relation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "position_ids = torch.tensor(position_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task_ids = torch.Tensor([1 for i in range(len(input_ids))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from multitask_format import MultiTaskModel, Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = home + '<name of your model folder>'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attach_task = Task(id = 0, name = 'attach prediction', type = \"seq_classification\", num_labels=2)\n",
    "relation_task = Task(id = 1, name = 'relation prediction', type = \"seq_classification\", num_labels = num_labels)\n",
    "tasks = [attach_task, relation_task]\n",
    "\n",
    "model = MultiTaskModel('bert-base-cased', tasks, len(tokenizer))\n",
    "output_model = model_path + '<name of your multitask .pth file output>'\n",
    "# output_model = model_path + 'multitask_stac.pth'\n",
    "print(output_model)\n",
    "checkpoint = torch.load(output_model, map_location='cuda')\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "model.to(device)\n",
    "print('loaded')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prediction on linear predicted attachments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "data_path = home + '<name of your linear preds pickle file>'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(data_path, 'rb') as f:\n",
    "    test_pred = pickle.load(f)\n",
    "\n",
    "input_ids, labels, raw = input_format(test_data, 10, relations=True, attach_preds=test_pred)\n",
    "# input_ids, labels, raw = input_format(test_data, 10, relations=False, attach_preds=test_pred)\n",
    "batch_tokenized = tokenizer(input_ids, return_tensors=\"pt\", padding=True, truncation=True, add_special_tokens=True)\n",
    "input_ids = batch_tokenized[\"input_ids\"].to(device) # list of token ids of dialogs in batch\n",
    "attention_masks = batch_tokenized[\"attention_mask\"].to(device)\n",
    "token_type_ids = batch_tokenized[\"token_type_ids\"].to(device)\n",
    "\n",
    "position_ids = position_ids_compute(tokenizer, input_ids, raw, labels)\n",
    "position_ids = torch.tensor(position_ids)\n",
    "\n",
    "task_ids = torch.Tensor([1 for i in range(len(input_ids))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(test_pred) == len(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset, DataLoader, SequentialSampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_data = TensorDataset(input_ids, attention_masks, token_type_ids, position_ids, task_ids)\n",
    "prediction_sampler = SequentialSampler(prediction_data)\n",
    "prediction_dataloader = DataLoader(prediction_data, sampler=prediction_sampler, batch_size=32)\n",
    "\n",
    "\n",
    "model.eval()\n",
    "\n",
    "predictions , true_labels = [], []\n",
    "\n",
    "for batch in prediction_dataloader:\n",
    "    batch = tuple(t.to(device) for t in batch)\n",
    "\n",
    "    b_input_ids, b_input_mask, b_token_types, b_position_ids, b_task_ids = batch\n",
    "\n",
    "    with torch.no_grad():\n",
    "        outputs, embed = model(b_input_ids,\n",
    "                     token_type_ids=b_token_types,\n",
    "                     attention_mask=b_input_mask,\n",
    "                     position_ids = b_position_ids,\n",
    "                     task_ids = b_task_ids)\n",
    "    logits = outputs[0]\n",
    "    logits = logits.detach().cpu().numpy()\n",
    "\n",
    "    predictions.append(logits)\n",
    "\n",
    "print('    DONE.')\n",
    "\n",
    "flat_prediction = np.concatenate(predictions, axis=0)\n",
    "flat_predictions = np.argmax(flat_prediction, axis=1).flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#add predictions to test pred attachments\n",
    "#make sure to keep all gold relations, even if not in predicted\n",
    "preds = []\n",
    "gold = []\n",
    "i = 0\n",
    "for n, g in enumerate(test_pred):\n",
    "  pred_tmp = []\n",
    "  gold_tmp = []\n",
    "  for rel in test_data[n]['relations']:\n",
    "    # if [rel['x'], rel['y']] in g:\n",
    "    #   gold_tmp.append([rel['x'], rel['y'], rel['type']])\n",
    "    if (rel['y']-rel['x']) <=10:\n",
    "      gold_tmp.append([rel['x'], rel['y'], rel['type']])\n",
    "  for p in g:\n",
    "    f = flat_predictions[i]\n",
    "    i += 1\n",
    "    pred_tmp.append([p[0], p[1], f])\n",
    "  gold.append(gold_tmp)\n",
    "  preds.append(pred_tmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#now preds is a list of lists of all the predicted relations on predicted attachments\n",
    "#and gold is a list of lists of all *gold* relations on predicted attachments\n",
    "len(preds), len(gold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#now we have to put them together so that in one place we have\n",
    "# dialogue index | x | y | gold relation (16 if not there)|pred relation\n",
    "comparisons = []\n",
    "for game in list(range(len(gold))):\n",
    "  gold_count = 0\n",
    "  goldgame = gold[game]\n",
    "  predgame = preds[game]\n",
    "  true_pos = [g for g in predgame if g in goldgame]\n",
    "  gold_count += len(true_pos)\n",
    "  rem_gold = [r for r in goldgame if r not in true_pos]\n",
    "  rem_pred = [r for r in predgame if r not in true_pos]\n",
    "  assert(len(goldgame) == len(true_pos) + len(rem_gold))\n",
    "  assert(len(predgame) == len(true_pos) + len(rem_pred))\n",
    "  for a in true_pos:\n",
    "    comparisons.append([game, a[0], a[1], a[2], a[2]])\n",
    "  #now decide for FPs and FNs whether they share a set of endpoints\n",
    "  rem_dict = defaultdict(list)\n",
    "  for rg in rem_gold: #false neg\n",
    "    rem_dict[(rg[0], rg[1])].append(('g', rg[2]))\n",
    "  for rp in rem_pred: #false pos\n",
    "    rem_dict[(rp[0], rp[1])].append(('p', rp[2]))\n",
    "\n",
    "  for it in rem_dict.keys():\n",
    "    p = 16\n",
    "    t = 16\n",
    "    for re in rem_dict[it]:\n",
    "      if re[0] == 'p':\n",
    "        p = re[1]\n",
    "      if re[0] == 'g':\n",
    "        t = re[1]\n",
    "        gold_count += 1\n",
    "    comparisons.append([game, it[0], it[1], t, p])\n",
    "  \n",
    "  assert(gold_count == (len(goldgame)))\n",
    "      \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#save multitask output \n",
    "with open(home + '<name of your pickle folder>/<name of your multitask preds pickle file>', 'wb') as f:\n",
    "    pickle.dump(comparisons, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import classification_report, ConfusionMatrixDisplay, confusion_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#all comparisons\n",
    "correct = [i[3] for i in comparisons]\n",
    "predicted = [i[4] for i in comparisons]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corr_all = [reverse_relations[i[3]] for i in comparisons]\n",
    "pred_all = [reverse_relations[i[4]] for i in comparisons]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(classification_report(corr_all,pred_all))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cm = confusion_matrix(correct,predicted)\n",
    "ConfusionMatrixDisplay(cm).plot()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bert_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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}