test_data_generator.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = './Wiki-Hades/valid.txt'\n",
    "import json\n",
    "# parse the json data into a list of dictionaries\n",
    "data = []\n",
    "with open(data_path, 'r') as f:\n",
    "    for line in f:\n",
    "        data.append(json.loads(line))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "masked_data = []\n",
    "for datapoint in data:\n",
    "    replaced_sent = datapoint['replaced']\n",
    "    lower_id = datapoint['replaced_ids'][0]\n",
    "    upper_id = datapoint['replaced_ids'][1]\n",
    "    replaced_sent = replaced_sent.split()\n",
    "    for i in range(lower_id, upper_id+1):\n",
    "        replaced_sent[i] = '[MASK]'\n",
    "    masked_sent = ' '.join(replaced_sent)\n",
    "    masked_data.append({'sentence': masked_sent, 'label': datapoint['hallucination'], 'replaced_ids': datapoint['replaced_ids']})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['bert.pooler.dense.weight', 'cls.seq_relationship.bias', 'cls.seq_relationship.weight', 'bert.pooler.dense.bias']\n",
      "- This IS expected if you are initializing BertForMaskedLM 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 BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "# load a bertformaskedlm model\n",
    "from transformers import BertForMaskedLM, AutoTokenizer\n",
    "import torch\n",
    "model = BertForMaskedLM.from_pretrained('bert-base-uncased').to('cuda')\n",
    "tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "100\n",
      "200\n",
      "300\n",
      "400\n",
      "500\n",
      "600\n",
      "700\n",
      "800\n",
      "900\n"
     ]
    }
   ],
   "source": [
    "# for each sentence in the masked data, predict the top 10 words that could fill in the mask`sa`\n",
    "\n",
    "labeled_multisample_data = []\n",
    "counter = 0\n",
    "for datapoint in masked_data:\n",
    "    input_ids = tokenizer.encode(datapoint['sentence'], return_tensors='pt').to('cuda')\n",
    "    mask_id = torch.where(input_ids == tokenizer.mask_token_id)[1].to('cuda')\n",
    "    predictions = model(input_ids)[0][:,mask_id,:].topk(10).indices.tolist()[0]\n",
    "    sentences = []\n",
    "    for i in range(10):\n",
    "        sentence = datapoint['sentence'].split()\n",
    "        for j in range(len(predictions)):\n",
    "            sentence[datapoint['replaced_ids'][0] + j] = tokenizer.decode(predictions[j][i])\n",
    "        sentences.append(' '.join(sentence))\n",
    "    labeled_multisample_data.append({'original':datapoint['sentence'], 'replaced_ids': datapoint['replaced_ids'],'sentences' : sentences, 'label': datapoint['label']})\n",
    "\n",
    "    if counter % 100 == 0:\n",
    "        print(counter)\n",
    "    counter += 1\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "100\n",
      "200\n",
      "300\n",
      "400\n",
      "500\n",
      "600\n",
      "700\n",
      "800\n",
      "900\n",
      "1000\n",
      "1100\n",
      "1200\n",
      "1300\n",
      "1400\n",
      "1500\n",
      "1600\n",
      "1700\n",
      "1800\n",
      "1900\n",
      "2000\n",
      "2100\n",
      "2200\n",
      "2300\n",
      "2400\n",
      "2500\n",
      "2600\n",
      "2700\n",
      "2800\n",
      "2900\n",
      "3000\n",
      "3100\n",
      "3200\n",
      "3300\n",
      "3400\n",
      "3500\n",
      "3600\n",
      "3700\n",
      "3800\n",
      "3900\n",
      "4000\n",
      "4100\n",
      "4200\n",
      "4300\n",
      "4400\n",
      "4500\n",
      "4600\n",
      "4700\n",
      "4800\n",
      "4900\n",
      "5000\n",
      "5100\n",
      "5200\n",
      "5300\n",
      "5400\n",
      "5500\n",
      "5600\n",
      "5700\n",
      "5800\n",
      "5900\n",
      "6000\n",
      "6100\n",
      "6200\n",
      "6300\n",
      "6400\n",
      "6500\n",
      "6600\n",
      "6700\n",
      "6800\n",
      "6900\n",
      "7000\n",
      "7100\n",
      "7200\n",
      "7300\n",
      "7400\n",
      "7500\n",
      "7600\n",
      "7700\n",
      "7800\n",
      "7900\n",
      "8000\n",
      "8100\n",
      "8200\n",
      "8300\n",
      "8400\n",
      "8500\n",
      "8600\n",
      "8700\n"
     ]
    }
   ],
   "source": [
    "# # for each sentence in the masked data, predict the top 20 words that could fill in the mask`sa`\n",
    "\n",
    "# labeled_multisample_data2 = []\n",
    "# counter = 0\n",
    "# for datapoint in masked_data:\n",
    "#     input_ids = tokenizer.encode(datapoint['sentence'], return_tensors='pt').to('cuda')\n",
    "#     mask_id = torch.where(input_ids == tokenizer.mask_token_id)[1].to('cuda')\n",
    "#     predictions = model(input_ids)[0][:,mask_id,:].topk(20).indices.tolist()[0]\n",
    "#     sentences = []\n",
    "#     for i in range(10):\n",
    "#         sentence = datapoint['sentence'].split()\n",
    "#         for j in range(len(predictions)):\n",
    "#             sentence[datapoint['replaced_ids'][0] + j] = tokenizer.decode(predictions[j][i])\n",
    "#         sentences.append(' '.join(sentence))\n",
    "#     labeled_multisample_data2.append({'original':datapoint['sentence'], 'replaced_ids': datapoint['replaced_ids'],'sentences' : sentences, 'label': datapoint['label']})\n",
    "\n",
    "#     if counter % 100 == 0:\n",
    "#         print(counter)\n",
    "#     counter += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pickle\n",
    "# with open('hades_aug_10_samples.pkl', 'rb') as f:\n",
    "#     labeled_multisample_data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of RobertaModel were not initialized from the model checkpoint at FacebookAI/roberta-base and are newly initialized: ['roberta.pooler.dense.bias', 'roberta.pooler.dense.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "from transformers import RobertaModel, RobertaTokenizer, AutoTokenizer\n",
    "import numpy as np\n",
    "embedding_model = RobertaModel.from_pretrained(\"FacebookAI/roberta-base\").to('cuda:2')\n",
    "embedding_tokenizer = AutoTokenizer.from_pretrained(\"FacebookAI/roberta-base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def top_k_variances(sentences , k=10, model=embedding_model, tokenizer=embedding_tokenizer):\n",
    "    embeddings = []\n",
    "    for sentence in sentences:\n",
    "        tokens = ['[CLS]'] + tokenizer.tokenize(sentence) + ['[SEP]']\n",
    "        if 50 > len(tokens):\n",
    "            tokens = tokens + ['[PAD]'] * (50 - len(tokens))\n",
    "        else:\n",
    "            tokens = tokens[:200]\n",
    "        attention_mask = [1 if i != '[PAD]' else 0 for i in tokens]\n",
    "        \n",
    "        token_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "        token_ids = torch.tensor(token_ids).unsqueeze(0)\n",
    "        attention_mask = torch.tensor(attention_mask).unsqueeze(0)\n",
    "        with torch.no_grad():\n",
    "            output = model(token_ids.to(\"cuda:2\"), attention_mask=attention_mask.to(\"cuda:2\"))\n",
    "        final_reps = output.last_hidden_state.detach().cpu()\n",
    "        cls_rep = final_reps[0][0].unsqueeze(0)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            embeddings.append(cls_rep)\n",
    "    \n",
    "    embeddings = np.array(embeddings, dtype=np.float32)\n",
    "    topk = np.var(embeddings, axis=0).squeeze()\n",
    "    topk = np.sort(topk)[-k:]\n",
    "    return topk[::-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    }
   ],
   "source": [
    "# utils.py from HADES\n",
    "from torch import nn\n",
    "import torch\n",
    "from torch.nn import functional as F\n",
    "import codecs\n",
    "import json\n",
    "import spacy\n",
    "from sklearn.metrics import classification_report, accuracy_score, hamming_loss, \\\n",
    "    f1_score, precision_score, recall_score, average_precision_score, roc_auc_score, confusion_matrix, \\\n",
    "    brier_score_loss\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def binary_eval(predy, testy, verbose=True, return_f1=False, predscore=None):\n",
    "    acc = accuracy_score(testy, predy)\n",
    "    f1 = f1_score(testy, predy, average=None)\n",
    "    precision = precision_score(testy, predy, average=None)\n",
    "    recall = recall_score(testy, predy, average=None)\n",
    "    epsilon = 1e-8\n",
    "\n",
    "    htn, hfp, hfn, htp = confusion_matrix(testy, predy).ravel()\n",
    "    hsensi = htp / (htp + hfn + epsilon)\n",
    "    hspec = htn / (hfp + htn + epsilon)\n",
    "    gmean = np.sqrt(hsensi*hspec)\n",
    "\n",
    "\n",
    "    info = \"Acc : {}\\nf1 : {}\\nprecision : {}\\nrecall : {}\\nG-mean : {}\".format(acc,\n",
    "            \" \".join([str(x) for x in f1]), \" \".join([str(x) for x in precision]),\n",
    "            \" \".join([str(x) for x in recall]), gmean)\n",
    "\n",
    "    if predscore is not None:\n",
    "        bss = brier_score_loss(testy, predscore)\n",
    "        roc_auc = roc_auc_score(testy, predscore)\n",
    "        info += \"\\nbss : {}\\nROC-AUC : {}\".format(bss, roc_auc)\n",
    "\n",
    "    if verbose:\n",
    "        print(info)\n",
    "\n",
    "    if return_f1:\n",
    "        return acc, f1, precision, recall, gmean, bss, roc_auc, info\n",
    "    else:\n",
    "        return acc, info\n",
    "\n",
    "\n",
    "def subsets(nums):\n",
    "    \"\"\"\n",
    "    :type nums: List[int]\n",
    "    :rtype: List[List[int]]\n",
    "    \"\"\"\n",
    "    ans = []\n",
    "    def dfs(curpos, tmp):\n",
    "        if tmp:\n",
    "            ans.append(tmp[:])\n",
    "        for i in range(curpos, len(nums)):\n",
    "            tmp.append(nums[i])\n",
    "            dfs(i+1, tmp)\n",
    "            tmp.pop(-1)\n",
    "    dfs(0, [])\n",
    "    return ans\n",
    "\n",
    "\n",
    "def sent_ner_bounds(sen, nlp=None):\n",
    "    if nlp is None:\n",
    "        nlp = spacy.load('en')\n",
    "    tokens, tags = [], []\n",
    "    print(sen)\n",
    "    for doc in nlp.pipe([sen]):\n",
    "        for token in doc:\n",
    "            tags.append(token.ent_iob_)\n",
    "            tokens.append(str(token))\n",
    "\n",
    "    rep_pos = []\n",
    "    vis = [False for _ in range(len(tags))]\n",
    "    for idx, tag in enumerate(tags):\n",
    "        if tag == 'O':\n",
    "            rep_pos.append([idx, idx])\n",
    "            vis[idx] = True\n",
    "        elif tag == 'B':\n",
    "            end = idx\n",
    "            for j in range(idx+1, len(tags)):\n",
    "                if tags[j] == 'I':\n",
    "                    end = j\n",
    "                else:\n",
    "                    break\n",
    "            rep_pos.append([idx, end])\n",
    "        elif tag == 'I':\n",
    "            continue\n",
    "\n",
    "    return ' '.join(tokens), rep_pos\n",
    "\n",
    "\n",
    "def remove_marked_sen(sen, start_id, end_id):\n",
    "    tokens = sen if type(sen) == list else sen.strip().split()\n",
    "    if tokens[start_id].startswith(\"===\") and tokens[end_id].endswith(\"===\"):\n",
    "        tokens[start_id] = tokens[start_id][3:]\n",
    "        tokens[end_id] = tokens[end_id][:-3]\n",
    "    return tokens\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\ntrainpath=\"../data_collections/wiki5k+sequential+topk_10+temp_1+context_1+rep_0.6+sample_1.annotate.finish\"\\n# trainpath=\"../data_collections/wiki5k+sequential+topk_10+temp_1+context_1+rep_0.6+sample_1.annotate\"\\ndp = DataProcessor()\\ntrain_examples = dp.get_examples(trainpath)\\nprint(len(train_examples))\\nprint(train_examples[2].sen)\\nprint(train_examples[2].idxs)\\nprint(train_examples[2].guid)\\nprint(dp.get_label_dist())\\n\\ntokenizer = BertTokenizer.from_pretrained(\"bert-base-uncased\")\\ntrain_dataset = HalluDataset(train_examples,tokenizer, 512, \"bert\")\\n\\ntrain_dataloader = data.DataLoader(dataset=train_dataset,\\n                                   batch_size=32,\\n                                   shuffle=False,\\n                                   num_workers=4,\\n                                   collate_fn=HalluDataset.pad)\\n\\nfor step, batch in enumerate(train_dataloader):\\n    input_ids, input_mask, segment_ids, predict_mask, label_ids, guids = batch\\n    print(\"id {} mask {} segment {}, pmask {}, label {}\".format(input_ids.size(), input_mask.size(),\\n                                                                segment_ids.size(), predict_mask.size(), label_ids.size()))\\n    print(\"guid {}\".format(\" \".join([str(guid) for guid in guids])))\\n'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# dataloader.py from HADES\n",
    "from torch.utils import data\n",
    "import torch\n",
    "import numpy as np\n",
    "from tqdm import tqdm, trange\n",
    "import collections\n",
    "import codecs\n",
    "import json\n",
    "\n",
    "\n",
    "class InputExample(object):\n",
    "\n",
    "    def __init__(self, guid, sen, idxs, label):\n",
    "        self.guid = guid\n",
    "        self.sen = sen\n",
    "        self.idxs = idxs\n",
    "        self.label = label\n",
    "\n",
    "\n",
    "class InputFeatures(object):\n",
    "\n",
    "    def __init__(self, guid, input_ids, input_mask, segment_ids,  predict_mask, label_id):\n",
    "        self.guid = guid\n",
    "        self.input_ids = input_ids\n",
    "        self.input_mask = input_mask\n",
    "        self.segment_ids = segment_ids\n",
    "        self.predict_mask = predict_mask\n",
    "        self.label_id = label_id\n",
    "\n",
    "\n",
    "class DataProcessor(object):\n",
    "    def __init__(self):\n",
    "        self.num_consist = 0\n",
    "        self.num_hallu = 0\n",
    "\n",
    "    def _read_data(self, input_file, require_uidx=False):\n",
    "        with open(input_file) as f:\n",
    "            # out_lines = []\n",
    "            out_lists = []\n",
    "            entries = f.read().strip().split(\"\\n\")\n",
    "            for entry in entries:\n",
    "                example = json.loads(entry.strip())\n",
    "                if \"hallucination\" not in example:\n",
    "                    label = -1\n",
    "                else:\n",
    "                    label = example[\"hallucination\"]\n",
    "                    if label not in [0, 1]:\n",
    "                        continue\n",
    "                if require_uidx:\n",
    "                    sen, token_ids, uidx = example[\"replaced\"], example[\"replaced_ids\"], example[\"idx\"]\n",
    "                    out_lists.append([sen, token_ids, label, uidx])\n",
    "                else:\n",
    "                    sen, token_ids = example[\"replaced\"], example[\"replaced_ids\"]\n",
    "                    out_lists.append([sen, token_ids, label])\n",
    "        return out_lists\n",
    "\n",
    "    def _create_examples(self, all_lists):\n",
    "        examples = []\n",
    "        for (i, one_lists) in enumerate(all_lists):\n",
    "            guid = i\n",
    "            if len(one_lists) == 3:  # Don't contain key \"idx\" in json file\n",
    "                sen, token_ids, label = one_lists\n",
    "            elif len(one_lists) == 4:  # Contain key \"idx\" in json file\n",
    "                sen, token_ids, label, guid = one_lists\n",
    "            else:\n",
    "                assert len(one_lists) == 3 or len(one_lists) == 4\n",
    "\n",
    "            if label == 0:\n",
    "                self.num_consist += 1\n",
    "            elif label == 1:\n",
    "                self.num_hallu += 1\n",
    "\n",
    "            examples.append(InputExample(\n",
    "                guid=guid, sen=sen, idxs=token_ids, label=label))\n",
    "        return examples\n",
    "\n",
    "    def get_examples(self, path, require_uidx=False):\n",
    "        return self._create_examples(\n",
    "            self._read_data(path, require_uidx))\n",
    "\n",
    "    def get_label_dist(self):\n",
    "        return [self.num_consist, self.num_hallu]\n",
    "\n",
    "def truncate(rep_subtokens, predict_mask, max_seq_length, rep_start_id, rep_end_id, mode=\"offline\"):\n",
    "    '''\n",
    "    Truncate the sequence if given a fixed context window. For example, given the following input sentence:\n",
    "    \"he signed a professional contract and promoted to the ===senior team=== where he managed to play for almost 3 years .\"\n",
    "    if the context window length is set as 4, the function will truncate the input as follows:\n",
    "\n",
    "    online mode: \"and promoted to the ===senior team===\"\n",
    "    offline mode: \"to the ===senior team=== where he\"\n",
    "    '''\n",
    "    if mode == \"offline\":\n",
    "        if len(rep_subtokens) > max_seq_length - 2:\n",
    "            mid_pt = int((rep_start_id + rep_end_id) / 2)\n",
    "            left_seq_length = int(max_seq_length / 2)\n",
    "            right_seq_length = max_seq_length - left_seq_length\n",
    "            if mid_pt - left_seq_length >= 0 and mid_pt + right_seq_length < len(rep_subtokens):\n",
    "                left_pt = mid_pt - left_seq_length\n",
    "                right_pt = mid_pt + right_seq_length\n",
    "            elif mid_pt - left_seq_length < 0 and mid_pt + right_seq_length < len(rep_subtokens):\n",
    "                left_pt = 0\n",
    "                right_pt = max_seq_length\n",
    "            elif mid_pt - left_seq_length >= 0 and mid_pt + right_seq_length >= len(rep_subtokens):\n",
    "                right_pt = len(rep_subtokens)\n",
    "                left_pt = len(rep_subtokens) - max_seq_length\n",
    "            elif mid_pt - left_seq_length < 0 and mid_pt + right_seq_length >= len(rep_subtokens):\n",
    "                left_pt = 0\n",
    "                right_pt = len(rep_subtokens)\n",
    "            rep_subtokens = rep_subtokens[left_pt:right_pt - 1]\n",
    "            predict_mask = predict_mask[left_pt:right_pt - 1]\n",
    "    else: # online\n",
    "        left_pt, right_pt = 0, rep_end_id + 1\n",
    "        if right_pt > max_seq_length - 2:\n",
    "            left_pt = right_pt - (max_seq_length - 2)\n",
    "        rep_subtokens = rep_subtokens[left_pt:right_pt]\n",
    "        predict_mask = predict_mask[left_pt:right_pt]\n",
    "    return rep_subtokens, predict_mask\n",
    "\n",
    "\n",
    "def example2feature(example, tokenizer, max_seq_length, model_name, mode=\"offline\"):\n",
    "    rep_start_id, rep_end_id = example.idxs\n",
    "    rep_tokens = remove_marked_sen(example.sen, rep_start_id, rep_end_id)\n",
    "\n",
    "    if 'xlnet' in model_name.lower():\n",
    "        rep_subtokens = []\n",
    "        predict_mask = []\n",
    "\n",
    "        for id, rep_token in enumerate(rep_tokens):\n",
    "            rep_subtoken = tokenizer.tokenize(rep_token)\n",
    "            if id >= rep_start_id and id <= rep_end_id:\n",
    "                rep_subtokens.extend(rep_subtoken)\n",
    "                predict_mask.extend(len(rep_subtoken) * [1])\n",
    "            else:\n",
    "                rep_subtokens.extend(rep_subtoken)\n",
    "                predict_mask.extend(len(rep_subtoken) * [0])\n",
    "\n",
    "        rep_subtokens, predict_mask = truncate(rep_subtokens, predict_mask, max_seq_length, rep_start_id, rep_end_id, mode=mode)\n",
    "\n",
    "        rep_subtokens.extend([\"<sep>\", \"<cls>\"])\n",
    "        predict_mask.extend([0, 0])\n",
    "\n",
    "    elif 'gpt' not in model_name.lower():\n",
    "\n",
    "        rep_subtokens = []\n",
    "        predict_mask = []\n",
    "        for id, rep_token in enumerate(rep_tokens):\n",
    "            rep_subtoken = tokenizer.tokenize(rep_token)\n",
    "            if id >= rep_start_id and id <= rep_end_id:\n",
    "                rep_subtokens.extend(rep_subtoken)\n",
    "                predict_mask.extend(len(rep_subtoken) * [1])\n",
    "            else:\n",
    "                rep_subtokens.extend(rep_subtoken)\n",
    "                predict_mask.extend(len(rep_subtoken) * [0])\n",
    "\n",
    "        rep_subtokens, predict_mask = truncate(rep_subtokens, predict_mask, max_seq_length, rep_start_id, rep_end_id, mode=mode)\n",
    "\n",
    "        rep_subtokens.insert(0, \"[CLS]\")\n",
    "        predict_mask.insert(0, 0)\n",
    "        rep_subtokens.append('[SEP]')\n",
    "        predict_mask.append(0)\n",
    "\n",
    "    elif 'gpt' in model_name.lower():\n",
    "        rep_subtokens = []\n",
    "        predict_mask = []\n",
    "\n",
    "        for id, rep_token in enumerate(rep_tokens):\n",
    "            rep_token = \" \"+rep_token if id!=0 else rep_token\n",
    "            rep_subtoken = tokenizer.tokenize(rep_token)\n",
    "            if id >= rep_start_id and id <= rep_end_id:\n",
    "                rep_subtokens.extend(rep_subtoken)\n",
    "                predict_mask.extend(len(rep_subtoken) * [1])\n",
    "            else:\n",
    "                rep_subtokens.extend(rep_subtoken)\n",
    "                predict_mask.extend(len(rep_subtoken) * [0])\n",
    "\n",
    "        rep_subtokens, predict_mask = truncate(rep_subtokens, predict_mask, max_seq_length, rep_start_id, rep_end_id, mode=mode)\n",
    "\n",
    "    input_ids = tokenizer.convert_tokens_to_ids(rep_subtokens)\n",
    "    segment_ids = [0] * len(input_ids)\n",
    "    input_mask = [1] * len(input_ids)\n",
    "\n",
    "    feat=InputFeatures(\n",
    "                guid=example.guid,\n",
    "                # tokens=tokens,\n",
    "                input_ids=input_ids,\n",
    "                input_mask=input_mask,\n",
    "                segment_ids=segment_ids,\n",
    "                predict_mask=predict_mask,\n",
    "                label_id=example.label)\n",
    "    return feat\n",
    "\n",
    "def get_examples_from_sen_tuple(sen, rep_pos):\n",
    "    examples = []\n",
    "    for uid, pos in enumerate(rep_pos):\n",
    "        examples.append(InputExample(guid=uid, sen=sen, idxs=pos, label=0))\n",
    "    return examples\n",
    "\n",
    "class HalluDataset(data.Dataset):\n",
    "    def __init__(self, examples, tokenizer, max_seq_length, model_name, task_mode=\"offline\"):\n",
    "        self.examples=examples\n",
    "        self.tokenizer=tokenizer\n",
    "        self.max_seq_length=max_seq_length\n",
    "        self.model_name = model_name\n",
    "        self.task_mode = task_mode\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.examples)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        feat=example2feature(self.examples[idx], self.tokenizer, self.max_seq_length,\n",
    "                             self.model_name, self.task_mode)\n",
    "        return feat.input_ids, feat.input_mask, feat.segment_ids, feat.predict_mask, feat.label_id, feat.guid\n",
    "\n",
    "    @classmethod\n",
    "    def pad(cls, batch):\n",
    "\n",
    "        seqlen_list = [len(sample[0]) for sample in batch]\n",
    "        maxlen = np.array(seqlen_list).max()\n",
    "\n",
    "        f = lambda x, seqlen: [sample[x] + [0] * (seqlen - len(sample[x])) for sample in batch] # 0: X for padding\n",
    "        input_ids_list = torch.LongTensor(f(0, maxlen))\n",
    "        input_mask_list = torch.LongTensor(f(1, maxlen))\n",
    "        segment_ids_list = torch.LongTensor(f(2, maxlen))\n",
    "        predict_mask_list = torch.ByteTensor(f(3, maxlen))\n",
    "        label_id = torch.LongTensor([sample[4] for sample in batch])\n",
    "        guids = [sample[5] for sample in batch]\n",
    "\n",
    "        return input_ids_list, input_mask_list, segment_ids_list, predict_mask_list, label_id, guids\n",
    "\n",
    "'''\n",
    "trainpath=\"../data_collections/wiki5k+sequential+topk_10+temp_1+context_1+rep_0.6+sample_1.annotate.finish\"\n",
    "# trainpath=\"../data_collections/wiki5k+sequential+topk_10+temp_1+context_1+rep_0.6+sample_1.annotate\"\n",
    "dp = DataProcessor()\n",
    "train_examples = dp.get_examples(trainpath)\n",
    "print(len(train_examples))\n",
    "print(train_examples[2].sen)\n",
    "print(train_examples[2].idxs)\n",
    "print(train_examples[2].guid)\n",
    "print(dp.get_label_dist())\n",
    "\n",
    "tokenizer = BertTokenizer.from_pretrained(\"bert-base-uncased\")\n",
    "train_dataset = HalluDataset(train_examples,tokenizer, 512, \"bert\")\n",
    "\n",
    "train_dataloader = data.DataLoader(dataset=train_dataset,\n",
    "                                   batch_size=32,\n",
    "                                   shuffle=False,\n",
    "                                   num_workers=4,\n",
    "                                   collate_fn=HalluDataset.pad)\n",
    "\n",
    "for step, batch in enumerate(train_dataloader):\n",
    "    input_ids, input_mask, segment_ids, predict_mask, label_ids, guids = batch\n",
    "    print(\"id {} mask {} segment {}, pmask {}, label {}\".format(input_ids.size(), input_mask.size(),\n",
    "                                                                segment_ids.size(), predict_mask.size(), label_ids.size()))\n",
    "    print(\"guid {}\".format(\" \".join([str(guid) for guid in guids])))\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertModel\n",
    "hades_model = BertModel.from_pretrained('google-bert/bert-base-uncased').to('cuda:2')\n",
    "hades_tokenizer = AutoTokenizer.from_pretrained('google-bert/bert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils import data\n",
    "test_datapath = './Wiki-Hades/valid.txt'\n",
    "dp = DataProcessor()\n",
    "test_examples = dp.get_examples(test_datapath)\n",
    "test_dataset = HalluDataset(test_examples, hades_tokenizer, 200, 'bert-base-uncased', \"offline\")\n",
    "test_dataloader = data.DataLoader(test_dataset, batch_size=1, num_workers=1, collate_fn=HalluDataset.pad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "Stored 0 vectors\n",
      "Stored 100 vectors\n",
      "Stored 200 vectors\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[14], line 16\u001b[0m\n\u001b[1;32m     14\u001b[0m datapoint \u001b[38;5;241m=\u001b[39m labeled_multisample_data[i]\n\u001b[1;32m     15\u001b[0m sentences \u001b[38;5;241m=\u001b[39m datapoint[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124msentences\u001b[39m\u001b[38;5;124m'\u001b[39m]\n\u001b[0;32m---> 16\u001b[0m topk_vars \u001b[38;5;241m=\u001b[39m \u001b[43mtop_k_variances\u001b[49m\u001b[43m(\u001b[49m\u001b[43msentences\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     18\u001b[0m train_x \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mconcatenate([cls_rep, topk_vars])[np\u001b[38;5;241m.\u001b[39mnewaxis, :]\n\u001b[1;32m     19\u001b[0m vector_reps\u001b[38;5;241m.\u001b[39mappend({\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mrepresentation\u001b[39m\u001b[38;5;124m'\u001b[39m: train_x, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlabel\u001b[39m\u001b[38;5;124m'\u001b[39m: datapoint[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlabel\u001b[39m\u001b[38;5;124m'\u001b[39m]})\n",
      "Cell \u001b[0;32mIn[7], line 9\u001b[0m, in \u001b[0;36mtop_k_variances\u001b[0;34m(sentences, k, model, tokenizer)\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m      8\u001b[0m     tokens \u001b[38;5;241m=\u001b[39m tokens[:\u001b[38;5;241m200\u001b[39m]\n\u001b[0;32m----> 9\u001b[0m attention_mask \u001b[38;5;241m=\u001b[39m [\u001b[38;5;241m1\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m i \u001b[38;5;241m!=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m[PAD]\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;241m0\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m tokens]\n\u001b[1;32m     11\u001b[0m token_ids \u001b[38;5;241m=\u001b[39m tokenizer\u001b[38;5;241m.\u001b[39mconvert_tokens_to_ids(tokens)\n\u001b[1;32m     12\u001b[0m token_ids \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mtensor(token_ids)\u001b[38;5;241m.\u001b[39munsqueeze(\u001b[38;5;241m0\u001b[39m)\n",
      "Cell \u001b[0;32mIn[7], line 9\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m      8\u001b[0m     tokens \u001b[38;5;241m=\u001b[39m tokens[:\u001b[38;5;241m200\u001b[39m]\n\u001b[0;32m----> 9\u001b[0m attention_mask \u001b[38;5;241m=\u001b[39m [\u001b[38;5;241m1\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[43mi\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m!=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43m[PAD]\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;241m0\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m tokens]\n\u001b[1;32m     11\u001b[0m token_ids \u001b[38;5;241m=\u001b[39m tokenizer\u001b[38;5;241m.\u001b[39mconvert_tokens_to_ids(tokens)\n\u001b[1;32m     12\u001b[0m token_ids \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mtensor(token_ids)\u001b[38;5;241m.\u001b[39munsqueeze(\u001b[38;5;241m0\u001b[39m)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# currently using bert-base-uncased for reps\n",
    "# for other models, reinitialize the model variable\n",
    "\n",
    "vector_reps = []\n",
    "\n",
    "for i, sample in enumerate(test_dataloader):\n",
    "    sample = tuple(t for t in sample[:-1])\n",
    "    input_ids, input_mask, _, _, _= sample\n",
    "    with torch.no_grad():\n",
    "        outputs = hades_model(input_ids.to(\"cuda:2\"), input_mask.to(\"cuda:2\"))\n",
    "    cls_rep = outputs.last_hidden_state[0][0]\n",
    "    cls_rep = cls_rep.detach().cpu().numpy()\n",
    "\n",
    "    datapoint = labeled_multisample_data[i]\n",
    "    sentences = datapoint['sentences']\n",
    "    topk_vars = top_k_variances(sentences)\n",
    "    \n",
    "    train_x = np.concatenate([cls_rep, topk_vars])[np.newaxis, :]\n",
    "    vector_reps.append({'representation': train_x, 'label': datapoint['label']})\n",
    "    if i % 100 == 0:\n",
    "        print(f\"Stored {i} vectors\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open('bert_validation_data.pkl', 'wb') as f:\n",
    "    pickle.dump(vector_reps, f)"
   ]
  }
 ],
 "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}