QAworkflow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "468dac70",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModel\n",
    "from copy import deepcopy\n",
    "import json\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import torch.nn as nn\n",
    "os.environ['TOKENIZERS_PARALLELISM'] = 'true'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8bc4cf75",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')\n",
    "train_examples = json.load(open('data/MetaQA_data/qa_train_1hop.txt.json', 'r'))\n",
    "valid_examples = json.load(open('data/MetaQA_data/qa_dev_1hop.txt.json', 'r'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b50ca46c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ent2idx = {}\n",
    "idx2ent = {}\n",
    "def get_entity_dict():\n",
    "    global ent2idx\n",
    "    global idx2ent\n",
    "    with open('data/MetaQA_data/entities.dict', 'r') as f:\n",
    "        lines = f.readlines()\n",
    "        for line in lines:\n",
    "            idx, ent = line.strip().split('\\t')\n",
    "            ent2idx[ent] = int(idx)\n",
    "            idx2ent[int(idx)] = ent\n",
    "get_entity_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bcac59f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class QA_Example:\n",
    "    def __init__(self, head, head_id, relation, tail, tail_id):\n",
    "        self.head = head\n",
    "        self.head_id = int(head_id)\n",
    "        self.question = relation\n",
    "        self.tail = tail\n",
    "        self.tail_id = [int(x) for x in tail_id]\n",
    "\n",
    "    def vectorize(self):\n",
    "        hq_tokenized = tokenizer(text=self.head,\n",
    "                                 text_pair=self.question,\n",
    "                                 max_length=48,\n",
    "                                 return_token_type_ids=True,\n",
    "                                 truncation=True)\n",
    "        \n",
    "        # t_tokenized = tokenizer(text=self.tail,\n",
    "        #                         max_length=10,\n",
    "        #                         return_token_type_ids=True,\n",
    "        #                         truncation=True)\n",
    "        \n",
    "        return {\n",
    "                # 'head_id': self.head_id,\n",
    "                'hq_token_ids': hq_tokenized['input_ids'],\n",
    "                'hq_token_type_ids': hq_tokenized['token_type_ids'],\n",
    "                'hq_mask': hq_tokenized['attention_mask'],\n",
    "                'all_answer_id': self.tail_id,\n",
    "                # 'all_tail_token_ids': t_tokenized['input_ids'],\n",
    "                # 'all_tail_token_type_ids': t_tokenized['token_type_ids'],\n",
    "                # 'all_tail_mask': t_tokenized['attention_mask']\n",
    "                }\n",
    "    \n",
    "class Ent_Example:\n",
    "    def __init__(self, entity):\n",
    "        self.entity = entity\n",
    "\n",
    "    def vectorize(self):\n",
    "        ent_tokenized = tokenizer(text=self.entity,\n",
    "                                max_length=16,\n",
    "                                return_token_type_ids=True,\n",
    "                                truncation=True)\n",
    "        \n",
    "        return {'ent_token_ids': ent_tokenized['input_ids'],\n",
    "                'ent_token_type_ids': ent_tokenized['token_type_ids'],\n",
    "                'ent_mask': ent_tokenized['attention_mask']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "636a9ccf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class QA_Dataset(torch.utils.data.dataset.Dataset):\n",
    "    def __init__(self, examples):\n",
    "        self.examples = []\n",
    "        for i in range(len(examples)):\n",
    "            self.examples.append(QA_Example(**examples[i]))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.examples)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return self.examples[index].vectorize()\n",
    "    \n",
    "class Ent_Dataset(torch.utils.data.dataset.Dataset):\n",
    "    def __init__(self, examples):\n",
    "        self.examples = []\n",
    "        for i in range(len(examples)):\n",
    "            self.examples.append(Ent_Example(examples[i]))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.examples)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return self.examples[index].vectorize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "66a15861",
   "metadata": {},
   "outputs": [],
   "source": [
    "def padding_batches(batch: list, pad_token_id, return_mask=False):\n",
    "    max_len = max([x.shape[0] for x in batch])\n",
    "    batch_size = len(batch)\n",
    "    token_batch = torch.LongTensor(batch_size, max_len).fill_(pad_token_id)\n",
    "    mask_batch = torch.zeros(batch_size, max_len)\n",
    "    for i, t in enumerate(batch):\n",
    "        token_batch[i, :len(t)].copy_(t)\n",
    "        if return_mask:\n",
    "            mask_batch[i, :len(t)].fill_(1)\n",
    "    if return_mask:\n",
    "        return token_batch, mask_batch\n",
    "    else:\n",
    "        return token_batch\n",
    "\n",
    "def _collate_fn(batch):\n",
    "    batch_dict = {}\n",
    "    batch_dict['hq_token_ids'], batch_dict['hq_mask'] = padding_batches([torch.LongTensor(x['hq_token_ids']) for x in batch],\n",
    "                                                                        tokenizer.pad_token_id,\n",
    "                                                                        return_mask=True)\n",
    "    batch_dict['hq_token_type_ids'] = padding_batches([torch.LongTensor(x['hq_token_type_ids']) for x in batch], 0)\n",
    "    batch_dict['all_answer_id'] = [x['all_answer_id'] for x in batch]\n",
    "    return batch_dict\n",
    "\n",
    "def _collate_fn_ent(batch):\n",
    "    batch_dict = {}\n",
    "    batch_dict['ent_token_ids'], batch_dict['ent_mask'] = padding_batches([torch.LongTensor(x['ent_token_ids']) for x in batch],\n",
    "                                                                          tokenizer.pad_token_id,\n",
    "                                                                          return_mask=True)\n",
    "    batch_dict['ent_token_type_ids'] = padding_batches([torch.LongTensor(x['ent_token_type_ids']) for x in batch], 0)\n",
    "    return batch_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f1e73497",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = QA_Dataset(train_examples)\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "                train_dataset,\n",
    "                batch_size=64,\n",
    "                shuffle=True,\n",
    "                collate_fn=_collate_fn,\n",
    "                num_workers=4)\n",
    "valid_dataset = QA_Dataset(valid_examples)\n",
    "valid_dataloader = torch.utils.data.DataLoader(\n",
    "                valid_dataset,\n",
    "                batch_size=64,\n",
    "                shuffle=True,\n",
    "                collate_fn=_collate_fn,\n",
    "                num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e91744f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_entity_loader():\n",
    "    ent_list = []\n",
    "    with open('data/MetaQA_data/entities.dict', 'r') as f:\n",
    "        lines = f.readlines()\n",
    "        for line in lines:\n",
    "            ent_list.append(line.strip().split('\\t')[1])\n",
    "    ent_dataset = Ent_Dataset(ent_list)\n",
    "    ent_dataloader = torch.utils.data.DataLoader(\n",
    "        ent_dataset,\n",
    "        batch_size=1024,\n",
    "        shuffle=False,\n",
    "        collate_fn=_collate_fn_ent,\n",
    "        num_workers=4)\n",
    "    return ent_dataloader\n",
    "ent_dataloader = get_entity_loader()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "69036ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def get_ent_embedding(model):\n",
    "    ent_embed = []\n",
    "    for ent_batch in ent_dataloader:\n",
    "        batch = move_to_cuda(ent_batch)\n",
    "        batch_embed = model.encode(model.ent_encoder,\n",
    "                                   input_ids=batch['ent_token_ids'],\n",
    "                                   token_type_ids=batch['ent_token_type_ids'],\n",
    "                                   attention_mask=batch['ent_mask'])\n",
    "        ent_embed.append(batch_embed)\n",
    "    return torch.cat(ent_embed)\n",
    "\n",
    "def move_to_cuda(sample):\n",
    "    if len(sample) == 0:\n",
    "        return {}\n",
    "\n",
    "    def _move_to_cuda(maybe_tensor):\n",
    "        if torch.is_tensor(maybe_tensor):\n",
    "            return maybe_tensor.cuda(non_blocking=True)\n",
    "        elif isinstance(maybe_tensor, dict):\n",
    "            return {key: _move_to_cuda(value) for key, value in maybe_tensor.items()}\n",
    "        elif isinstance(maybe_tensor, list):\n",
    "            return [_move_to_cuda(x) for x in maybe_tensor]\n",
    "        elif isinstance(maybe_tensor, tuple):\n",
    "            return [_move_to_cuda(x) for x in maybe_tensor]\n",
    "        else:\n",
    "            return maybe_tensor\n",
    "\n",
    "    return _move_to_cuda(sample)\n",
    "\n",
    "def get_one_hot_label(answer_ids, ent_num=43234, alpha_smooth=0):\n",
    "    label = torch.zeros(len(answer_ids), ent_num) + alpha_smooth / ent_num\n",
    "    for i, answers in enumerate(answer_ids):\n",
    "        for answer in answers:\n",
    "            label[i, int(answer)] = 1 - alpha_smooth\n",
    "    return label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a5ee8a63",
   "metadata": {},
   "outputs": [],
   "source": [
    "class QAModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.ent_encoder = AutoModel.from_pretrained('bert-base-uncased').cuda()\n",
    "        self.ent_encoder.load_state_dict(torch.load('checkpoint/MetaQA_2023-04-23-0538.24/epoch2tail_bert.mdl'))\n",
    "        self.qa_encoder = AutoModel.from_pretrained('bert-base-uncased').cuda()\n",
    "        self.qa_encoder.load_state_dict(torch.load('checkpoint/MetaQA_2023-04-23-0538.24/epoch2hr_bert.mdl'))\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "        self.ent_embed = get_ent_embedding(self)\n",
    "    \n",
    "    def encode(self, model, input_ids, token_type_ids, attention_mask):\n",
    "        embed = model(input_ids=input_ids,\n",
    "                      token_type_ids=token_type_ids,\n",
    "                      attention_mask=attention_mask,\n",
    "                      return_dict=True)\n",
    "        embed = embed.last_hidden_state\n",
    "        embed = self.convert_features(attention_mask, embed)\n",
    "        return embed\n",
    "\n",
    "    def forward(self, qa_batch_dict):\n",
    "        qa_embed = self.encode(self.qa_encoder,\n",
    "                               input_ids=qa_batch_dict['hq_token_ids'],\n",
    "                               token_type_ids=qa_batch_dict['hq_token_type_ids'],\n",
    "                               attention_mask=qa_batch_dict['hq_mask'])\n",
    "        \n",
    "        qa_embed = qa_embed[:, None, :]\n",
    "        score = (qa_embed * self.ent_embed).sum(dim=2)\n",
    "        return self.softmax(score)\n",
    "\n",
    "    def convert_features(self, mask, last_hidden_state, convert_type='mean'):\n",
    "        if convert_type=='first':\n",
    "            output_vector = last_hidden_state[:,0,:]\n",
    "        else:\n",
    "            input_mask_expanded = mask.unsqueeze(-1).expand(last_hidden_state.size()).float()\n",
    "            sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, 1)\n",
    "            sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-4)\n",
    "            output_vector = sum_embeddings / sum_mask\n",
    "        output_vector = nn.functional.normalize(output_vector, dim=1)\n",
    "        return output_vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "72e7482d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.seq_relationship.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.bias', 'cls.predictions.transform.dense.bias']\n",
      "- This IS expected if you are initializing BertModel 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 BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.seq_relationship.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.bias', 'cls.predictions.transform.dense.bias']\n",
      "- This IS expected if you are initializing BertModel 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 BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "model = QAModel().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b4d0bd20",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def evaluate_accuracy(model, dataloader):\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for i, batch in enumerate(tqdm(dataloader)):\n",
    "        pred_ans = model(move_to_cuda(batch)).argmax(dim=1)\n",
    "        label = batch['all_answer_id']\n",
    "        for j, ans in enumerate(pred_ans):\n",
    "            if ans in label[j]:\n",
    "                correct += 1\n",
    "            total += 1\n",
    "    return correct / total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "1998e718",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def evaluate_mrr(model, dataloader):\n",
    "    total = 0\n",
    "    n = 0\n",
    "    for i, batch in enumerate(tqdm(dataloader)):\n",
    "        pred_ans = model(move_to_cuda(batch))\n",
    "        label = batch['all_answer_id']\n",
    "        for j, ans in enumerate(label):\n",
    "            rank = min([(pred_ans[j]>=pred_ans[j, x]).sum() for x in ans])\n",
    "            total += 1 / rank\n",
    "            n += 1\n",
    "    return total / n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "3e826bef",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "acc:0.00, mrr:0.00, loss:0.0005, avg loss:0.000:   1%|          | 26/3266 [00:05<12:08,  4.45it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[35], line 15\u001b[0m\n\u001b[1;32m     13\u001b[0m optimizer\u001b[39m.\u001b[39mzero_grad()\n\u001b[1;32m     14\u001b[0m loss \u001b[39m=\u001b[39m loss_f(x, label\u001b[39m.\u001b[39mcuda())\n\u001b[0;32m---> 15\u001b[0m loss\u001b[39m.\u001b[39;49mbackward()\n\u001b[1;32m     16\u001b[0m total_loss \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m loss\u001b[39m.\u001b[39mitem()\n\u001b[1;32m     17\u001b[0m avg_loss \u001b[39m=\u001b[39m total_loss \u001b[39m/\u001b[39m (i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m)\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/_tensor.py:429\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    426\u001b[0m     \u001b[39m# All strings are unicode in Python 3.\u001b[39;00m\n\u001b[1;32m    427\u001b[0m     \u001b[39mreturn\u001b[39;00m torch\u001b[39m.\u001b[39m_tensor_str\u001b[39m.\u001b[39m_str(\u001b[39mself\u001b[39m, tensor_contents\u001b[39m=\u001b[39mtensor_contents)\n\u001b[0;32m--> 429\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mbackward\u001b[39m(\n\u001b[1;32m    430\u001b[0m     \u001b[39mself\u001b[39m, gradient\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m, retain_graph\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m, create_graph\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m, inputs\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m\n\u001b[1;32m    431\u001b[0m ):\n\u001b[1;32m    432\u001b[0m     \u001b[39mr\u001b[39m\u001b[39m\"\"\"Computes the gradient of current tensor w.r.t. graph leaves.\u001b[39;00m\n\u001b[1;32m    433\u001b[0m \n\u001b[1;32m    434\u001b[0m \u001b[39m    The graph is differentiated using the chain rule. If the tensor is\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    476\u001b[0m \u001b[39m            used to compute the attr::tensors.\u001b[39;00m\n\u001b[1;32m    477\u001b[0m \u001b[39m    \"\"\"\u001b[39;00m\n\u001b[1;32m    478\u001b[0m     \u001b[39mif\u001b[39;00m has_torch_function_unary(\u001b[39mself\u001b[39m):\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "loss_f = nn.BCELoss()\n",
    "optimizer = torch.optim.AdamW(model.qa_encoder.parameters(), lr=5e-5, weight_decay=1e-4)\n",
    "epochs = 1\n",
    "acc = 0\n",
    "mrr = 0\n",
    "for epoch in range(epochs):\n",
    "    model.train()\n",
    "    with tqdm(total=len(train_dataloader)) as pbar:\n",
    "        total_loss = 0\n",
    "        for i, batch in enumerate(train_dataloader):\n",
    "            x = model(move_to_cuda(batch))\n",
    "            label = get_one_hot_label(batch['all_answer_id'], alpha_smooth=0.0)\n",
    "            optimizer.zero_grad()\n",
    "            loss = loss_f(x, label.cuda())\n",
    "            loss.backward()\n",
    "            total_loss += loss.item()\n",
    "            avg_loss = total_loss / (i + 1)\n",
    "            optimizer.step()\n",
    "            pbar.set_description(f\"acc:{acc:.2f}, mrr:{mrr:.2f}, loss:{loss.item():.4f}, avg loss:{avg_loss:.3f}\")\n",
    "            pbar.update(1)\n",
    "    model.eval()\n",
    "    acc = evaluate_accuracy(model, valid_dataloader)\n",
    "    mrr = evaluate_mrr(model, valid_dataloader)\n",
    "print(f\"acc:{acc:.2f}\")\n",
    "torch.save(model.state_dict(), f'qa1hopmodel_acc{acc:.2f}')"
   ]
  }
 ],
 "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}