multi-token generation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "043dcb6a",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/katemarg/miniconda3/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "import sys\n",
    "import os\n",
    "import json\n",
    "import pprint\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import re\n",
    "import torch\n",
    "from collections import Counter\n",
    "from transformers import AutoModel, AutoTokenizer, pipeline\n",
    "from tqdm import tqdm\n",
    "from sklearn.metrics import f1_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "38b40238",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "BASE_DIR = \"/Users/katemarg/PycharmProjects/temporal_robustness_evaluation\"\n",
    "DATA_DIR = os.path.join(BASE_DIR, 'data')\n",
    "CKPT_DIR = os.path.join(BASE_DIR, 'pretrained_models')\n",
    "RES_DIR = os.path.join(BASE_DIR, 'results')\n",
    "LOG_DIR = os.path.join(BASE_DIR, 'logs')\n",
    "CACHE_DIR = os.path.join(BASE_DIR, 'cached')\n",
    "\n",
    "TEMPLAMA_ORIG_DIR = os.path.join(\"data\", \"templama\", \"test.json\")\n",
    "TEMPLAMA_REPRODUCED_DIR = os.path.join(\"templama_docker\", \"reproduce1\", \"templama\", \"test.jsonl\")\n",
    "TEMPLAMA_NEW_DIR = os.path.join(DATA_DIR, \"dynamic-templama\",\n",
    "                                \"dataset_from_2019-1-1_to_2022-12-31_per_quarter\", \n",
    "                                \"test.jsonl\")\n",
    "lm = \"cardiffnlp/twitter-roberta-base-mar2022\"\n",
    "# dataset_filepath=CACHE_DIR+'/{}_dynamic-templama_multiple_masks.pt'.format(lm)\n",
    "dataset_filepath=CACHE_DIR+'/cardiffnlp-twitter-roberta-base-mar2022_dynamic-templama_multiple_masks.pt'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ccc54a5",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f309acec",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "data_dict_multi_token = torch.load(dataset_filepath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "01767375",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Load pre-trained model tokenizer (vocabulary)\n",
    "tokenizer = AutoTokenizer.from_pretrained(lm, use_fast=False, add_prefix_space=True)\n",
    "\n",
    "CLS = tokenizer.cls_token\n",
    "PAD = tokenizer.pad_token\n",
    "SEP = tokenizer.sep_token\n",
    "MASK = tokenizer.mask_token\n",
    "\n",
    "mask_id = tokenizer.mask_token_id\n",
    "sep_id = tokenizer.sep_token_id\n",
    "cls_id = tokenizer.cls_token_id\n",
    "pad_id = tokenizer.pad_token_id\n",
    "\n",
    "special_ids = [mask_id, sep_id, cls_id, pad_id]\n",
    "def tokenizer_return_id(text, filter_special_tokens=False):\n",
    "    \"\"\"\n",
    "    Text to token ids for a string.\n",
    "    \"\"\"\n",
    "    output = tokenizer(text)\n",
    "    if filter_special_tokens:\n",
    "        token_ids = [i for i in output['input_ids'] if i not in tokenizer.all_special_ids]\n",
    "    else:\n",
    "        token_ids = [i for i in output['input_ids'] ]\n",
    "    return token_ids\n",
    "\n",
    "def tokenize_batch(batch):\n",
    "    \"\"\"\n",
    "    Text to token ids for a list of strings.\n",
    "    \"\"\"\n",
    "#     return [tokenizer.convert_tokens_to_ids(sent) for sent in batch]\n",
    "    return [tokenizer_return_id(sent) for sent in batch]\n",
    "\n",
    "def untokenize_id(ids):\n",
    "    \"\"\"\n",
    "    Token ids to strings.\n",
    "    \"\"\"\n",
    "#     return [tokenizer.convert_tokens_to_ids(sent) for sent in batch]\n",
    "    return [tokenizer.decode(_id) for _id in ids]\n",
    "\n",
    "\n",
    "def untokenize_batch(batch, filter_special_tokens=False):\n",
    "    \"\"\"\n",
    "    Token ids to strings for a list of ids.\n",
    "    \"\"\"\n",
    "#     return [tokenizer.convert_ids_to_tokens(sent) for sent in batch]\n",
    "#     print(label_id)\n",
    "    if filter_special_tokens:\n",
    "        _batch = []\n",
    "        for sent in batch:\n",
    "            _batch.append([x for x in sent if x not in special_ids])\n",
    "        batch = _batch\n",
    "#         return [tokenizer.decode(label_id) for label_id in batch if label_id not in special_ids]\n",
    "#     else:\n",
    "    return [untokenize_id(sent) for sent in batch]\n",
    "\n",
    "def detokenize(sent):\n",
    "    \"\"\" Roughly detokenizes (mainly undoes wordpiece) \"\"\"\n",
    "    # not sure what this does.... their code\n",
    "    new_sent = []\n",
    "    for i, tok in enumerate(sent):\n",
    "        if tok.startswith(\"##\"):\n",
    "            new_sent[len(new_sent) - 1] = new_sent[len(new_sent) - 1] + tok[2:]\n",
    "        else:\n",
    "            new_sent.append(tok)\n",
    "    return new_sent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0defa520",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def split_dataset(data):\n",
    "    \"\"\"\n",
    "    Split temporal dataset Dt to D_unchanged, D_new and D_updated compared to D_(t-1) for all t.\n",
    "    Specifically:\n",
    "    - D_unchanged: data where text_t = text_(t-1) & label_t = label_(t-1)\n",
    "    - D_updated: data where text_t = text_(t-1) & label_t != label_(t-1)\n",
    "    - D_new: data where text_t not in D_(t-1)\n",
    "    - D_deleted: data that exist in D_(t-1) but not in D_t\n",
    "\n",
    "    Args:\n",
    "        data: a dictionary with keys the time (year/quarter/month) and values dictionaries\n",
    "        data = {\n",
    "                '2019-Q1':\n",
    "                    {\n",
    "                    'text': [list of text],\n",
    "                    'labels': [list of labels],\n",
    "                    'labels_ids': [list of label token ids -- for a given model/tokenizer],\n",
    "                    'relations' [list of Wikidata relations]\n",
    "                    },\n",
    "                '2019-Q2': {...}\n",
    "                }\n",
    "\n",
    "    Returns:\n",
    "        D_unchanged, D_new, D_updated, D_deleted\n",
    "    \"\"\"\n",
    "    unchanged_t, new_t, updated_t, deleted_t = {}, {}, {}, {}\n",
    "    quarters = list(data.keys())\n",
    "    t_0 = quarters[0]  # t=t0\n",
    "    t_1 = quarters[0]  # t-1\n",
    "\n",
    "    for t in quarters[1:]:\n",
    "        if t in ['2022-Q3', '2022-Q4']:\n",
    "            continue # skip last two quarters of 2022\n",
    "        data_t = data[t]      # D_t\n",
    "        data_t_1 = data[t_1]  # D_(t-1)\n",
    "\n",
    "        unchanged_t[t] = {key: [] for key in data_t.keys()}\n",
    "        new_t[t] = {key: [] for key in data_t.keys()}\n",
    "        updated_t[t] = {key: [] for key in data_t.keys()}\n",
    "        deleted_t[t] = {key: [] for key in data_t.keys()}\n",
    "\n",
    "        for i in range(0, len(data_t['text'])):\n",
    "            text_t = data_t['text'][i]\n",
    "            labels_ids_t = data_t['labels_ids'][i][0]\n",
    "            if text_t in data_t_1['text']:\n",
    "                t_1_index = data_t_1['text'].index(text_t)\n",
    "                labels_inds_t_1 = data_t_1['labels_ids'][t_1_index][0]\n",
    "                if labels_ids_t == labels_inds_t_1:\n",
    "                    # text_t = text_t-1 & label_t = label_t-1\n",
    "                    # add to D_unchanged\n",
    "                    for key in data_t.keys():\n",
    "                        unchanged_t[t][key].append(data_t[key][i])\n",
    "                else:\n",
    "                    # text_t = text_(t-1) & label_t != label_(t-1)\n",
    "                    # add to D_updated\n",
    "                    for key in ['text', 'relation']:\n",
    "                        updated_t[t][key].append(data_t[key][i])\n",
    "            else:\n",
    "                # text_t not in D_(t-1) texts\n",
    "                # add to D_new\n",
    "                for key in data_t.keys():\n",
    "                    new_t[t][key].append(data_t[key][i])\n",
    "        for j in range(0, len(data_t_1['text'])):\n",
    "            text_t_1 = data_t_1['text'][j]\n",
    "            if text_t_1 not in data_t['text']:\n",
    "                # text_(t+1) not in D_t\n",
    "                # add to D_deleted\n",
    "                for key in data_t_1.keys():\n",
    "                    deleted_t[t][key].append(data_t_1[key][j])\n",
    "        t_1 = t\n",
    "\n",
    "        assert len(data_t['text']) == len(unchanged_t[t]['text']) + len(updated_t[t]['text']) + len(new_t[t]['text'])\n",
    "        print(\n",
    "            't={}: From total {} samples in D_t, {} are unchanged, {} are updated, {} are deleted and {} are new, compared to D_(t-1).'.format(\n",
    "                t,\n",
    "                len(data_t['text']),\n",
    "                len(unchanged_t[t]['text']),\n",
    "                len(updated_t[t]['text']),\n",
    "                len(deleted_t[t]['text']),\n",
    "                len(new_t[t]['text'])))\n",
    "    return unchanged_t, new_t, updated_t, deleted_t, data[t_0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ea7f3c5",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Split dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5ec4914b",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t=2019-Q2: From total 7962 samples in D_t, 7728 are unchanged, 36 are updated, 206 are deleted and 198 are new, compared to D_(t-1).\n",
      "t=2019-Q3: From total 8011 samples in D_t, 7633 are unchanged, 65 are updated, 264 are deleted and 313 are new, compared to D_(t-1).\n",
      "t=2019-Q4: From total 7989 samples in D_t, 7742 are unchanged, 51 are updated, 218 are deleted and 196 are new, compared to D_(t-1).\n",
      "t=2020-Q1: From total 7872 samples in D_t, 7359 are unchanged, 107 are updated, 524 are deleted and 406 are new, compared to D_(t-1).\n",
      "t=2020-Q2: From total 7848 samples in D_t, 7699 are unchanged, 30 are updated, 143 are deleted and 119 are new, compared to D_(t-1).\n",
      "t=2020-Q3: From total 7900 samples in D_t, 7666 are unchanged, 40 are updated, 142 are deleted and 194 are new, compared to D_(t-1).\n",
      "t=2020-Q4: From total 7892 samples in D_t, 7682 are unchanged, 26 are updated, 192 are deleted and 184 are new, compared to D_(t-1).\n",
      "t=2021-Q1: From total 7819 samples in D_t, 7381 are unchanged, 69 are updated, 442 are deleted and 369 are new, compared to D_(t-1).\n",
      "t=2021-Q2: From total 7771 samples in D_t, 7608 are unchanged, 32 are updated, 178 are deleted and 131 are new, compared to D_(t-1).\n",
      "t=2021-Q3: From total 7764 samples in D_t, 7545 are unchanged, 32 are updated, 195 are deleted and 187 are new, compared to D_(t-1).\n",
      "t=2021-Q4: From total 7758 samples in D_t, 7594 are unchanged, 31 are updated, 140 are deleted and 133 are new, compared to D_(t-1).\n",
      "t=2022-Q1: From total 7689 samples in D_t, 7416 are unchanged, 47 are updated, 294 are deleted and 226 are new, compared to D_(t-1).\n",
      "t=2022-Q2: From total 7686 samples in D_t, 7598 are unchanged, 21 are updated, 70 are deleted and 67 are new, compared to D_(t-1).\n"
     ]
    }
   ],
   "source": [
    "# Split dataset\n",
    "unchanged_t, new_t, updated_t, deleted_t, orig = split_dataset(data_dict_multi_token)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a84fa9cc",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b71665a8",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "lm = \"cardiffnlp/twitter-roberta-base-mar2022\"\n",
    "fill_mask_model = pipeline(\n",
    "    'fill-mask', model=lm, framework=\"pt\",\n",
    "    tokenizer=tokenizer, top_k=100\n",
    ")\n",
    "model = fill_mask_model.model\n",
    "model.eval()\n",
    "cuda = torch.cuda.is_available()\n",
    "if cuda:\n",
    "    model = model.cuda()\n",
    "year='2019-Q1'\n",
    "batch_size = 100\n",
    "N = batch_size # number of shots\n",
    "temperature = 1.0\n",
    "burnin = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ee89579c",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def generate_step(out, gen_idx, temperature=None, top_k=0, sample=False, return_list=True, num_masks=0):\n",
    "    \"\"\" Generate a word from out[gen_idx]\n",
    "\n",
    "    args:\n",
    "        - out (torch.Tensor): tensor of logits of size batch_size x seq_len x vocab_size\n",
    "        - gen_idx (int): location for which to generate for\n",
    "        - top_k (int): if >0, only sample from the top k most probable words\n",
    "        - sample (Bool): if True, sample from full distribution. Overridden by top_k\n",
    "    \"\"\"\n",
    "    logits = out[:, gen_idx]\n",
    "    if num_masks==1:  # if single-token then return the argmax\n",
    "        idx = torch.argmax(logits, dim=-1)\n",
    "        return idx.tolist() if return_list else idx\n",
    "    else:\n",
    "        if temperature is not None:\n",
    "            logits = logits / temperature\n",
    "        if top_k > 0:\n",
    "            kth_vals, kth_idx = logits.topk(top_k, dim=-1)\n",
    "            dist = torch.distributions.categorical.Categorical(logits=kth_vals)\n",
    "            idx = kth_idx.gather(dim=1, index=dist.sample().unsqueeze(-1)).squeeze(-1)\n",
    "        elif sample:\n",
    "            dist = torch.distributions.categorical.Categorical(logits=logits)\n",
    "            idx = dist.sample().squeeze(-1)\n",
    "        else:\n",
    "            idx = torch.argmax(logits, dim=-1)\n",
    "        return idx.tolist() if return_list else idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e8c81cc0",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "text_list = data_dict_multi_token[year]['text']\n",
    "labels_list = data_dict_multi_token[year]['labels']\n",
    "labels_ids_list = data_dict_multi_token[year]['labels_ids']\n",
    "relation_list = data_dict_multi_token[year]['relation']\n",
    "num_answers_list = data_dict_multi_token[year]['num_answers']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "662977e1",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## For a single test example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 374,
   "id": "e274c8ed",
   "metadata": {
    "scrolled": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Example: Lionel Messi plays for <mask> <mask>., Labels: [' FC Barcelona'], Ids: [5429, 4612]\n"
     ]
    }
   ],
   "source": [
    "M = 5 # max masks to try all from 1, ..., M\n",
    "i = 0\n",
    "text_i = text_list[i][0]\n",
    "# text_i = ['Lionel M plays for <mask> <mask>.'][0]\n",
    "labels_i = labels_list[i][0]\n",
    "labels_ids_i = labels_ids_list[i][0] # there is an 'extra' list that is why we put [0]\n",
    "relation_i = relation_list[i]\n",
    "num_answers_i = num_answers_list[i]\n",
    "print('Example: {}, Labels: {}, Ids: {}'.format(text_i, labels_i, labels_ids_i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 375,
   "id": "1385c107",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# tokenize text\n",
    "tokenized_sentence_orig = tokenizer_return_id(text_i)\n",
    "\n",
    "# mask indices to generate\n",
    "mask_inds_orig = list(np.where(np.array(tokenized_sentence_orig) == mask_id)[0])\n",
    "\n",
    "# number of masks (tokens)\n",
    "gold_num_masks = len(mask_inds_orig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 376,
   "id": "4d284d45",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# find first and last masks\n",
    "first_mask_idx = mask_inds_orig[0]\n",
    "last_mask_idx = mask_inds_orig[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 377,
   "id": "e38caf6b",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# split tokenized sentence to list of token ids before and after the masks\n",
    "before_mask_ids = tokenized_sentence_orig[:first_mask_idx]\n",
    "after_mask_ids = tokenized_sentence_orig[last_mask_idx+1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 378,
   "id": "131551d1",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "assert mask_id not in before_mask_ids\n",
    "assert mask_id not in after_mask_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 379,
   "id": "1db783ee",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 15350, 9711, 1974, 13, 50264, 50264, 4, 2]"
      ]
     },
     "execution_count": 379,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_sentence_orig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 380,
   "id": "0d637bf0",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# create all M combinations of masks [1,M]\n",
    "all_M_mask_combos = []\n",
    "for i in range (1,M+1):\n",
    "    \n",
    "    all_M_mask_combos.append([mask_id for _ in range(i)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 381,
   "id": "5ada1faf",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[50264],\n",
       " [50264, 50264],\n",
       " [50264, 50264, 50264],\n",
       " [50264, 50264, 50264, 50264],\n",
       " [50264, 50264, 50264, 50264, 50264]]"
      ]
     },
     "execution_count": 381,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_M_mask_combos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 382,
   "id": "d3bb3aea",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# add all M mask combos to list of token ids\n",
    "tokenized_sentence_list = []\n",
    "for mask_seq in all_M_mask_combos:\n",
    "    tokenized_sentence_list.append(before_mask_ids + mask_seq + after_mask_ids) # list of M lists (variable number of masks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 383,
   "id": "e41abfaf",
   "metadata": {
    "scrolled": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "######################\n",
      "We try with 1  mask(s)\n",
      "######################\n",
      "[0, 15350, 9711, 1974, 13, 50264, 4, 2]\n",
      "0:  Barcelona, 14.143179893493652\n",
      "1:  Barcelona, 14.143179893493652\n",
      "2:  Barcelona, 14.143179893493652\n",
      "3:  Barcelona, 14.143179893493652\n",
      "4:  Barcelona, 14.143179893493652\n",
      "5:  Barcelona, 14.143179893493652\n",
      "6:  Barcelona, 14.143179893493652\n",
      "7:  Barcelona, 14.143179893493652\n",
      "8:  Barcelona, 14.143179893493652\n",
      "9:  Barcelona, 14.143179893493652\n",
      "10:  Barcelona, 14.143179893493652\n",
      "11:  Barcelona, 14.143179893493652\n",
      "12:  Barcelona, 14.143179893493652\n",
      "13:  Barcelona, 14.143179893493652\n",
      "14:  Barcelona, 14.143179893493652\n",
      "15:  Barcelona, 14.143179893493652\n",
      "16:  Barcelona, 14.143179893493652\n",
      "17:  Barcelona, 14.143179893493652\n",
      "18:  Barcelona, 14.143179893493652\n",
      "19:  Barcelona, 14.143179893493652\n",
      "######################\n",
      "We try with 2  mask(s)\n",
      "######################\n",
      "[0, 15350, 9711, 1974, 13, 50264, 50264, 4, 2]\n",
      "0:  Lionel Messi, 10.740882396697998\n",
      "1:  Lionel Messi, 10.740882396697998\n",
      "2:  Lionel Messi, 10.740882396697998\n",
      "3:  Lionel Messi, 10.740882396697998\n",
      "4:  Lionel Messi, 10.740882396697998\n",
      "5:  Lionel Messi, 10.740882396697998\n",
      "6:  Lionel Messi, 10.740882396697998\n",
      "7:  Lionel Messi, 10.740882396697998\n",
      "8:  Lionel Messi, 10.740882396697998\n",
      "9:  Lionel Messi, 10.740882396697998\n",
      "10:  Lionel Messi, 10.740882396697998\n",
      "11:  Lionel Messi, 10.740882396697998\n",
      "12:  Lionel Messi, 10.740882396697998\n",
      "13:  Lionel Messi, 10.740882396697998\n",
      "14:  Lionel Messi, 10.740882396697998\n",
      "15:  Lionel Messi, 10.740882396697998\n",
      "16:  Lionel Messi, 10.740882396697998\n",
      "17:  Lionel Messi, 10.740882396697998\n",
      "18:  Lionel Messi, 10.740882396697998\n",
      "19:  Lionel Messi, 10.740882396697998\n",
      "######################\n",
      "We try with 3  mask(s)\n",
      "######################\n",
      "[0, 15350, 9711, 1974, 13, 50264, 50264, 50264, 4, 2]\n",
      "0:  Messi's club, 10.54391606648763\n",
      "1:  Messi's team, 10.43660799662272\n",
      "2:  Messi's team, 10.43660799662272\n",
      "3:  the Barcelona team, 10.18064022064209\n",
      "4:  the Lionel Messi, 9.912489891052246\n",
      "5:  the Lionel Messi, 9.912489891052246\n",
      "6:  the Lionel Messi, 9.912489891052246\n",
      "7:  the Lionel Messi, 9.912489891052246\n",
      "8:  the Lionel Messi, 9.912489891052246\n",
      "9:  the Lionel Messi, 9.912489891052246\n",
      "10:  the Lionel Messi, 9.912489891052246\n",
      "11:  the Lionel Messi, 9.912489891052246\n",
      "12:  the Lionel Messi, 9.912489891052246\n",
      "13:  the Lionel Messi, 9.912489891052246\n",
      "14:  the Lionel Messi, 9.912489891052246\n",
      "15:  the Lionel Messi, 9.912489891052246\n",
      "16:  Lionel Messi club, 9.753942330678305\n",
      "17:  Lionel Messi club, 9.753942330678305\n",
      "18:  Lionel Messi club, 9.753942330678305\n",
      "19:  Messi' Club, 9.543644587198893\n",
      "######################\n",
      "We try with 4  mask(s)\n",
      "######################\n",
      "[0, 15350, 9711, 1974, 13, 50264, 50264, 50264, 50264, 4, 2]\n",
      "0:  a club in Europe, 8.914562225341797\n",
      "1:  a country in Africa, 8.656488418579102\n",
      "2:  a club in Spain, 8.557990550994873\n",
      "3:  the club of England, 8.516459703445435\n",
      "4:  the club of Portugal, 8.389087915420532\n",
      "5:  the team not Messi, 8.325096130371094\n",
      "6:  the club of Messi, 8.306910872459412\n",
      "7:  the club of Barcelona, 8.177204132080078\n",
      "8:  the club of Barcelona, 8.177204132080078\n",
      "9:  the club of Barcelona, 8.177204132080078\n",
      "10:  the club of Barcelona, 8.177204132080078\n",
      "11:  the world of football, 8.1375972032547\n",
      "12:  the world of football, 8.1375972032547\n",
      "13:  a club called Barcelona, 8.083427429199219\n",
      "14:  the champions of Europe, 8.052391409873962\n",
      "15:  the champions of Europe, 8.052391409873962\n",
      "16:  a club like Barcelona, 8.044433355331421\n",
      "17:  a team like Barcelona, 8.009209632873535\n",
      "18:  the United of Portugal, 7.821687459945679\n",
      "19:  the team you support, 7.519671082496643\n",
      "######################\n",
      "We try with 5  mask(s)\n",
      "######################\n",
      "[0, 15350, 9711, 1974, 13, 50264, 50264, 50264, 50264, 50264, 4, 2]\n",
      "0:  the Messi-Lad, 8.192081356048584\n",
      "1:  the Barcelona Barca team, 6.792348384857178\n",
      "2:  the club of Lionel Messi, 6.583143854141236\n",
      "3:  the Barcelona Barca Messi, 6.563374996185303\n",
      "4:  Lionel Messi. Lionel Messi, 6.54958209991455\n",
      "5:  the Barcelona Barcelona Barcelona Barcelona, 6.510838651657105\n",
      "6:  the Barcelona Barcelona FC team, 6.363993120193482\n",
      "7:  the world's biggest club, 5.9062609553337095\n",
      "8:  the world and is Messi, 5.821213743090629\n",
      "9:  the Messi's Club Barcelona, 5.726065373420715\n",
      "10:  the Barcelona Madrid National team, 5.386617413163185\n",
      "11:  Lionel Messi. Not Messi, 5.220057857036591\n",
      "12:  the Barcelona United, Barcelona, 5.05019382238388\n",
      "13:  the world of Manchester United, 4.96308644413948\n",
      "14:  Lionel Messi at Messi Camp, 4.959240126609802\n",
      "15:  the club that is Messi, 4.7834427744150165\n",
      "16:  the club of Barca, 4.7003414034843445\n",
      "17:  Manchester United and Manchester City, 4.689992141723633\n",
      "18:  Manchester City and Manchester United, 4.657993966341019\n",
      "19:  Barça and not Barcelona, 4.480504083633423\n"
     ]
    }
   ],
   "source": [
    "# We do not know a priori the correct number of masks so we try all in range 1, ..., M (M=5)\n",
    "# tokenized_sentence_b_np, tokenized_sentence_b, logits_m, logits_m = [], [], [], []\n",
    "top_ranked_gen_token_seqs = []\n",
    "top_ranked_log_probs = []\n",
    "for num_mask_j in range(M):\n",
    "    generated_seq_tokens = [[] for _ in range(batch_size)]  # for each trial\n",
    "    logits_list = [[] for _ in range(batch_size)]\n",
    "    \n",
    "    print('##'*11)\n",
    "    print('We try with {}  mask(s)'.format(num_mask_j+1))\n",
    "    print('##'*11)\n",
    "    tokenized_sentence = tokenized_sentence_list[num_mask_j]\n",
    "    print(tokenized_sentence)\n",
    "    # list of indeces of the mask tokens\n",
    "    mask_inds = list(np.where(np.array(tokenized_sentence) == mask_id)[0])\n",
    "    \n",
    "    # create batch (same input, N times)\n",
    "    tokenized_sentence_b = [tokenized_sentence for _ in range(batch_size)]\n",
    "    \n",
    "    for m in mask_inds:\n",
    "        # tensor\n",
    "        inp = torch.tensor(tokenized_sentence_b).cuda() if cuda else torch.tensor(tokenized_sentence_b)\n",
    "\n",
    "        # get logits\n",
    "        out = model(inp)\n",
    "        logits = out.logits  # batch_size x max_len x vocab\n",
    "\n",
    "        # get new ids\n",
    "        idxs_b = generate_step(logits, gen_idx=m, top_k=10, temperature=temperature, # sample=(m < burnin)\n",
    "                              )\n",
    "\n",
    "        # replace mask with predicted token id\n",
    "        tokenized_sentence_b_np = np.array(tokenized_sentence_b)                       \n",
    "        tokenized_sentence_b_np[:,m] = np.array(idxs_b)\n",
    "        tokenized_sentence_b = tokenized_sentence_b_np.tolist() \n",
    "\n",
    "        # the following code does not work *and I don't know why*\n",
    "        # for jj in range(len(idxs_b)):\n",
    "        #     print('before: {}, predicted token id: {}'.format(tokenized_sentence_b[jj][m],idxs_b[jj]))\n",
    "        #     tokenized_sentence_b[jj][m] = idxs_b[jj]\n",
    "\n",
    "        assert sorted(idxs_b) == sorted([sent[m] for sent in tokenized_sentence_b])\n",
    "\n",
    "        # find logits\n",
    "        logits_m = logits[:,m]  # logits for the mask in position m\n",
    "        logits_b = logits_m[:,idxs_b].tolist()[0]  # logits for sampled tokens\n",
    "\n",
    "        assert len(idxs_b) == len(logits_b)\n",
    "\n",
    "        # add generated tokens and corresponding logits to lists\n",
    "        for j in range(N):\n",
    "            generated_seq_tokens[j].append(idxs_b[j])\n",
    "            logits_list[j].append(logits_b[j])\n",
    "    \n",
    "    # calculate sum of logits for each generated sequence of tokens\n",
    "    sum_logits = np.array(logits_list).sum(axis=-1) / len(mask_inds)\n",
    "\n",
    "    # finding ranking (return indices of the parallel lists for logits and generated tokens)\n",
    "    ranked_inds_of_list = sum_logits.argsort()[::-1]\n",
    "\n",
    "    # ranked logits sum in descending order\n",
    "    ranked_logits = sum_logits[ranked_inds_of_list]\n",
    "\n",
    "    # ranked generated tokens in descending order\n",
    "    ranked_generated_tokens = np.array(generated_seq_tokens)[ranked_inds_of_list]\n",
    "    \n",
    "    for no, p in enumerate(ranked_generated_tokens[:20]):\n",
    "        print('{}: {}, {}'.format(no,\"\".join(untokenize_id(p)), ranked_logits[no]))\n",
    "    \n",
    "    top_ranked_gen_token_seqs.append(ranked_generated_tokens.tolist()[0])\n",
    "    top_ranked_log_probs.append(ranked_logits.tolist()[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 396,
   "id": "619ee1f0",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[4612],\n",
       " [15350, 9711],\n",
       " [9711, 18, 950],\n",
       " [10, 950, 11, 1005],\n",
       " [5, 9711, 12, 574, 625]]"
      ]
     },
     "execution_count": 396,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_ranked_gen_token_seqs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 397,
   "id": "93db17ef",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[14.143179893493652,\n",
       " 10.740882396697998,\n",
       " 10.54391606648763,\n",
       " 8.914562225341797,\n",
       " 8.192081356048584]"
      ]
     },
     "execution_count": 397,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_ranked_log_probs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e73e1329",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 395,
   "id": "131623ad",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# compare with gold labels\n",
    "import evaluate\n",
    "# assert len(labels_ids_i) == M\n",
    "f1_micro_list, f1_macro_list = [], []\n",
    "rouge_list, bleu_list, bleu_uni_list, bert_score_list = [], [], [], []\n",
    "\n",
    "rouge = evaluate.load('rouge')\n",
    "bleu = evaluate.load(\"bleu\")\n",
    "bertscore = load(\"bertscore\")\n",
    "\n",
    "gold_ids = labels_ids_i\n",
    "gold_tok = labels_i  # list of strings\n",
    "\n",
    "for m in range(len(top_ranked_log_probs)):\n",
    "    pred_ids_m = top_ranked_gen_token_seqs[m]\n",
    "    pred_tok_m = [\"\".join(untokenize_id(top_ranked_gen_token_seqs[m]))] # list of strings\n",
    "    \n",
    "    # F1 score\n",
    "    # F1_micro calculates metrics globally by counting the total true positives, \n",
    "    # false negatives and false positives.\n",
    "    if len(gold_ids) == len(pred_ids_m):\n",
    "        f1_micro_list.append(f1_score(gold_ids,pred_ids_m, average='micro')) \n",
    "        # F1_macro calculates metrics for each label, and finds their unweighted mean. \n",
    "        # This does not take label imbalance into account.\n",
    "        f1_macro_list.append(f1_score(gold_ids,pred_ids_m, average='macro'))\n",
    "    else:\n",
    "        f1_micro_list.append(0.0)\n",
    "        f1_macro_list.append(0.0)\n",
    "        \n",
    "    # BLEU\n",
    "    bleu_list.append(bleu.compute(references=gold_tok,\n",
    "          predictions=pred_tok_m)['bleu'])\n",
    "    # unigrams\n",
    "    bleu_uni_list.append(bleu.compute(references=gold_tok,\n",
    "              predictions=pred_tok_m)['precisions'][0])\n",
    "\n",
    "    # ROUGE\n",
    "    rouge_list.append(rouge.compute(references=gold_tok,\n",
    "                  predictions=pred_tok_m, use_aggregator=True,\n",
    "                                    use_stemmer=True)['rouge1'].mid.fmeasure)\n",
    "    # BERT_SCORE\n",
    "    bert_score_list.append(bertscore.compute(references=gold_tok,\n",
    "              predictions=pred_tok_m, lang=\"en\")['f1'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 398,
   "id": "729730de",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The gold label is << FC Barcelona>> with 2 masks, while the predictions are:\n",
      "\n",
      "<< Barcelona>> with 1 mask(s)\n",
      "<< Lionel Messi>> with 2 mask(s)\n",
      "<< Messi's club>> with 3 mask(s)\n",
      "<< a club in Europe>> with 4 mask(s)\n",
      "<< the Messi-Lad>> with 5 mask(s)\n"
     ]
    }
   ],
   "source": [
    "pred_strings = [\"\".join(untokenize_id(pred)) for pred in top_ranked_gen_token_seqs]\n",
    "print('The gold label is <<{}>> with {} masks, while the predictions are:\\n'.format(gold_tok[0], gold_num_masks))\n",
    "for i in range(1,M+1):\n",
    "    print('<<{}>> with {} mask(s)'.format(pred_strings[i-1], i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 401,
   "id": "a6f72664",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1.0, 0.0, 0.0, 0.0, 0.0]"
      ]
     },
     "execution_count": 401,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bleu_uni_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 403,
   "id": "2ac1f1bf",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>masks</th>\n",
       "      <th>predictions</th>\n",
       "      <th>f1_micro</th>\n",
       "      <th>f1_macro</th>\n",
       "      <th>rouge</th>\n",
       "      <th>bleu</th>\n",
       "      <th>blue_uni_precision</th>\n",
       "      <th>bert_score</th>\n",
       "      <th>norm_log_probs</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>Barcelona</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.666667</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.955906</td>\n",
       "      <td>14.143180</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>Lionel Messi</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.889325</td>\n",
       "      <td>10.740882</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>Messi's club</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.877186</td>\n",
       "      <td>10.543916</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>a club in Europe</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.863004</td>\n",
       "      <td>8.914562</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>the Messi-Lad</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.842462</td>\n",
       "      <td>8.192081</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   masks        predictions  f1_micro  f1_macro     rouge  bleu  \\\n",
       "0      1          Barcelona       0.0       0.0  0.666667   0.0   \n",
       "1      2       Lionel Messi       0.0       0.0  0.000000   0.0   \n",
       "2      3       Messi's club       0.0       0.0  0.000000   0.0   \n",
       "3      4   a club in Europe       0.0       0.0  0.000000   0.0   \n",
       "4      5      the Messi-Lad       0.0       0.0  0.000000   0.0   \n",
       "\n",
       "   blue_uni_precision  bert_score  norm_log_probs  \n",
       "0                 1.0    0.955906       14.143180  \n",
       "1                 0.0    0.889325       10.740882  \n",
       "2                 0.0    0.877186       10.543916  \n",
       "3                 0.0    0.863004        8.914562  \n",
       "4                 0.0    0.842462        8.192081  "
      ]
     },
     "execution_count": 403,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(data={'masks':list(range(1,M+1)), 'predictions':pred_strings, \n",
    " 'f1_micro': f1_micro_list, 'f1_macro':f1_macro_list,\n",
    "'rouge': rouge_list, 'bleu': bleu_list, 'blue_uni_precision':bleu_uni_list, 'bert_score': bert_score_list,\n",
    "                  'norm_log_probs':top_ranked_log_probs\n",
    "                  })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e45403d",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f422678b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## I think we should consider every argmax generated token sequence for all M masks as equal predictions, and keep the highest score from any. In this case the first."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97d997bf",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 311,
   "id": "f13db2cb",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "n_samples = 5\n",
    "batch_size = 1\n",
    "max_len = 40\n",
    "top_k = 100\n",
    "temperature = 1.0\n",
    "generation_mode = \"parallel-sequential\"\n",
    "leed_out_len = 5 # max_len\n",
    "burnin = 250\n",
    "sample = True\n",
    "max_iter = 500\n",
    "\n",
    "# # Choose the prefix context\n",
    "# seed_text = \"[CLS]\".split()\n",
    "# bert_sents = generate(n_samples, seed_text=seed_text, batch_size=batch_size, max_len=max_len,\n",
    "#                       generation_mode=generation_mode,\n",
    "#                       sample=sample, top_k=top_k, temperature=temperature, burnin=burnin, max_iter=max_iter,\n",
    "#                       cuda=cuda)\n",
    "\n",
    "def sequential_generation(seed_text, batch_size=10, max_len=15, leed_out_len=15, \n",
    "                          top_k=0, temperature=None, sample=True, cuda=False):\n",
    "    \"\"\" Generate one word at a time, in L->R order \"\"\"\n",
    "    seed_len = len(seed_text)\n",
    "    batch = get_init_text(seed_text, max_len, batch_size)\n",
    "    \n",
    "    for ii in range(max_len):\n",
    "        inp = [sent[:seed_len+ii+leed_out_len]+[sep_id] for sent in batch]\n",
    "        inp = torch.tensor(batch).cuda() if cuda else torch.tensor(batch)\n",
    "        out = model(inp)\n",
    "        idxs = generate_step(out, gen_idx=seed_len+ii, top_k=top_k, temperature=temperature, sample=sample)\n",
    "        for jj in range(batch_size):\n",
    "            batch[jj][seed_len+ii] = idxs[jj]\n",
    "        \n",
    "    return untokenize_batch(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "609483e5",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[5429, 4612]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels_ids_i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "43a406a0",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 4868,   534],\n",
       "       [ 4868,   534],\n",
       "       [ 4868,   534],\n",
       "       [ 4868,   534],\n",
       "       [ 4868,   534],\n",
       "       [    5,  4612],\n",
       "       [    5,  4612],\n",
       "       [    5,  4612],\n",
       "       [ 4612,  4612],\n",
       "       [ 4612,  4612],\n",
       "       [ 2361,   412],\n",
       "       [ 2361,   412],\n",
       "       [ 2361,   412],\n",
       "       [ 2361,   412],\n",
       "       [    5,  3622],\n",
       "       [    5,   950],\n",
       "       [    5,   950],\n",
       "       [    5,   950],\n",
       "       [ 2361,   315],\n",
       "       [ 2361,   315],\n",
       "       [ 2361,   315],\n",
       "       [ 2361,   315],\n",
       "       [ 2361,   315],\n",
       "       [ 2361,   315],\n",
       "       [ 2361,   315],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [15350,  9711],\n",
       "       [    5,   165],\n",
       "       [ 1554,   412],\n",
       "       [ 1554,   412],\n",
       "       [ 1554,   412],\n",
       "       [ 1554,   412],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [ 1731,  3245],\n",
       "       [    5,  9711],\n",
       "       [    5,  9711],\n",
       "       [ 1554,   315],\n",
       "       [ 1554,   315],\n",
       "       [ 9711,  2009],\n",
       "       [ 9711,  1731],\n",
       "       [ 4612,   122],\n",
       "       [    5,  4174],\n",
       "       [    5,  6772],\n",
       "       [    5,  6772],\n",
       "       [ 1731, 19393],\n",
       "       [ 1731, 19393]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ranked_generated_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "90e5d2d9",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "for j in range(0,N):\n",
    "    gold_ids = labels_ids_i\n",
    "    pred_ids_j = ranked_generated_tokens[j]\n",
    "    gold_tok = labels_i  # list of strings\n",
    "    pred_tok_j = [\"\".join(untokenize_id(ranked_generated_tokens[j]))] # list of strings\n",
    "    \n",
    "    # F1 score\n",
    "    # F1_micro calculates metrics globally by counting the total true positives, \n",
    "    # false negatives and false positives.\n",
    "    f1_micro_list.append(f1_score(gold_ids,pred_ids_j, average='micro')) \n",
    "    # F1_macro calculates metrics for each label, and finds their unweighted mean. \n",
    "    # This does not take label imbalance into account.\n",
    "    f1_macro_list.append(f1_score(gold_ids,pred_ids_j, average='macro'))\n",
    "    \n",
    "    # BLEU\n",
    "    bleu_list.append(bleu.compute(references=gold_tok,\n",
    "              predictions=pred_tok_j)['bleu'])\n",
    "\n",
    "    # ROUGE\n",
    "    rouge_list.append(rouge.compute(references=gold_tok,\n",
    "                  predictions=pred_tok_j))\n",
    "    \n",
    "    # exact match / P@1\n",
    "max_f1_micro = np.array(f1_micro_list).max()\n",
    "max_f1_macro = np.array(f1_macro_list).max()\n",
    "max_bleu = np.array(bleu_list).max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "27d0a3ae",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.5 0.3333333333333333 0.0\n"
     ]
    }
   ],
   "source": [
    "print(max_f1_micro, max_f1_macro, max_bleu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dd1f2a9",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def generate_step(out, gen_idx, temperature=None, top_k=0, sample=False, return_list=True):\n",
    "    \"\"\" Generate a word from from out[gen_idx]\n",
    "    \n",
    "    args:\n",
    "        - out (torch.Tensor): tensor of logits of size batch_size x seq_len x vocab_size\n",
    "        - gen_idx (int): location for which to generate for\n",
    "        - top_k (int): if >0, only sample from the top k most probable words\n",
    "        - sample (Bool): if True, sample from full distribution. Overridden by top_k \n",
    "    \"\"\"\n",
    "    logits = out[:, gen_idx]\n",
    "    if temperature is not None:\n",
    "        logits = logits / temperature\n",
    "    if top_k > 0:\n",
    "        kth_vals, kth_idx = logits.topk(top_k, dim=-1)\n",
    "        dist = torch.distributions.categorical.Categorical(logits=kth_vals)\n",
    "        idx = kth_idx.gather(dim=1, index=dist.sample().unsqueeze(-1)).squeeze(-1)\n",
    "    elif sample:\n",
    "        dist = torch.distributions.categorical.Categorical(logits=logits)\n",
    "        idx = dist.sample().squeeze(-1)\n",
    "    else:\n",
    "        idx = torch.argmax(logits, dim=-1)\n",
    "    return idx.tolist() if return_list else idx\n",
    "  \n",
    "def get_init_text(seed_text, max_len, batch_size = 1, rand_init=False):\n",
    "    \"\"\" Get initial sentence by padding seed_text with either masks or random words to max_len \"\"\"\n",
    "    batch = [seed_text + [MASK] * max_len + [SEP] for _ in range(batch_size)]\n",
    "    #if rand_init:\n",
    "    #    for ii in range(max_len):\n",
    "    #        init_idx[seed_len+ii] = np.random.randint(0, len(tokenizer.vocab))\n",
    "    \n",
    "    return tokenize_batch(batch)\n",
    "\n",
    "def printer(sent, should_detokenize=True):\n",
    "    if should_detokenize:\n",
    "        sent = detokenize(sent)[1:-1]\n",
    "    print(\" \".join(sent))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb818693",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Generation modes as functions\n",
    "import math\n",
    "import time\n",
    "\n",
    "def parallel_sequential_generation(seed_text, batch_size=10, max_len=15, top_k=0, temperature=None, max_iter=300, burnin=200,\n",
    "                                   cuda=False, print_every=10, verbose=True):\n",
    "    \"\"\" Generate for one random position at a timestep\n",
    "    \n",
    "    args:\n",
    "        - burnin: during burn-in period, sample from full distribution; afterwards take argmax\n",
    "    \"\"\"\n",
    "    seed_len = len(seed_text)\n",
    "    batch = get_init_text(seed_text, max_len, batch_size)\n",
    "    \n",
    "    for ii in range(max_iter):\n",
    "        kk = np.random.randint(0, max_len)\n",
    "        for jj in range(batch_size):\n",
    "            batch[jj][seed_len+kk] = mask_id\n",
    "        inp = torch.tensor(batch).cuda() if cuda else torch.tensor(batch)\n",
    "        out = model(inp)\n",
    "        topk = top_k if (ii >= burnin) else 0\n",
    "        idxs = generate_step(out, gen_idx=seed_len+kk, top_k=topk, temperature=temperature, sample=(ii < burnin))\n",
    "        for jj in range(batch_size):\n",
    "            batch[jj][seed_len+kk] = idxs[jj]\n",
    "            \n",
    "        if verbose and np.mod(ii+1, print_every) == 0:\n",
    "            for_print = tokenizer.convert_ids_to_tokens(batch[0])\n",
    "            for_print = for_print[:seed_len+kk+1] + ['(*)'] + for_print[seed_len+kk+1:]\n",
    "            print(\"iter\", ii+1, \" \".join(for_print))\n",
    "            \n",
    "    return untokenize_batch(batch)\n",
    "\n",
    "def parallel_generation(seed_text, batch_size=10, max_len=15, top_k=0, temperature=None, max_iter=300, sample=True, \n",
    "                        cuda=False, print_every=10, verbose=True):\n",
    "    \"\"\" Generate for all positions at each time step \"\"\"\n",
    "    seed_len = len(seed_text)\n",
    "    batch = get_init_text(seed_text, max_len, batch_size)\n",
    "    \n",
    "    for ii in range(max_iter):\n",
    "        inp = torch.tensor(batch).cuda() if cuda else torch.tensor(batch)\n",
    "        out = model(inp)\n",
    "        for kk in range(max_len):\n",
    "            idxs = generate_step(out, gen_idx=seed_len+kk, top_k=top_k, temperature=temperature, sample=sample)\n",
    "            for jj in range(batch_size):\n",
    "                batch[jj][seed_len+kk] = idxs[jj]\n",
    "            \n",
    "        if verbose and np.mod(ii, print_every) == 0:\n",
    "            print(\"iter\", ii+1, \" \".join(tokenizer.convert_ids_to_tokens(batch[0])))\n",
    "    \n",
    "    return untokenize_batch(batch)\n",
    "            \n",
    "def sequential_generation(seed_text, batch_size=10, max_len=15, leed_out_len=15, \n",
    "                          top_k=0, temperature=None, sample=True, cuda=False):\n",
    "    \"\"\" Generate one word at a time, in L->R order \"\"\"\n",
    "    seed_len = len(seed_text)\n",
    "    batch = get_init_text(seed_text, max_len, batch_size)\n",
    "    \n",
    "    for ii in range(max_len):\n",
    "        inp = [sent[:seed_len+ii+leed_out_len]+[sep_id] for sent in batch]\n",
    "        inp = torch.tensor(batch).cuda() if cuda else torch.tensor(batch)\n",
    "        out = model(inp)\n",
    "        idxs = generate_step(out, gen_idx=seed_len+ii, top_k=top_k, temperature=temperature, sample=sample)\n",
    "        for jj in range(batch_size):\n",
    "            batch[jj][seed_len+ii] = idxs[jj]\n",
    "        \n",
    "    return untokenize_batch(batch)\n",
    "\n",
    "def generate(n_samples, seed_text=\"[CLS]\", batch_size=10, max_len=25, \n",
    "             generation_mode=\"parallel-sequential\",\n",
    "             sample=True, top_k=100, temperature=1.0, burnin=200, max_iter=500,\n",
    "             cuda=False, print_every=1):\n",
    "    # main generation function to call\n",
    "    sentences = []\n",
    "    n_batches = math.ceil(n_samples / batch_size)\n",
    "    start_time = time.time()\n",
    "    for batch_n in range(n_batches):\n",
    "        if generation_mode == \"parallel-sequential\":\n",
    "            batch = parallel_sequential_generation(seed_text, batch_size=batch_size, max_len=max_len, top_k=top_k,\n",
    "                                                   temperature=temperature, burnin=burnin, max_iter=max_iter, \n",
    "                                                   cuda=cuda, verbose=False)\n",
    "        elif generation_mode == \"sequential\":\n",
    "            batch = sequential_generation(seed_text, batch_size=batch_size, max_len=max_len, top_k=top_k, \n",
    "                                          temperature=temperature, leed_out_len=leed_out_len, sample=sample,\n",
    "                                          cuda=cuda)\n",
    "        elif generation_mode == \"parallel\":\n",
    "            batch = parallel_generation(seed_text, batch_size=batch_size,\n",
    "                                        max_len=max_len, top_k=top_k, temperature=temperature, \n",
    "                                        sample=sample, max_iter=max_iter, \n",
    "                                        cuda=cuda, verbose=False)\n",
    "        \n",
    "        if (batch_n + 1) % print_every == 0:\n",
    "            print(\"Finished batch %d in %.3fs\" % (batch_n + 1, time.time() - start_time))\n",
    "            start_time = time.time()\n",
    "        \n",
    "        sentences += batch\n",
    "    return sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2654e42",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "n_samples = 5\n",
    "batch_size = 5\n",
    "max_len = 40\n",
    "top_k = 100\n",
    "temperature = 1.0\n",
    "generation_mode = \"parallel-sequential\"\n",
    "leed_out_len = 5 # max_len\n",
    "burnin = 250\n",
    "sample = True\n",
    "max_iter = 500\n",
    "\n",
    "# Choose the prefix context\n",
    "seed_text = \"[CLS]\".split()\n",
    "bert_sents = generate(n_samples, seed_text=seed_text, batch_size=batch_size, max_len=max_len,\n",
    "                      generation_mode=generation_mode,\n",
    "                      sample=sample, top_k=top_k, temperature=temperature, burnin=burnin, max_iter=max_iter,\n",
    "                      cuda=cuda)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e059ebb",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "for sent in bert_sents:\n",
    "  printer(sent, should_detokenize=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae878754",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def tokenizer_return_id(text, filter_special_tokens=False):\n",
    "    \"\"\"\n",
    "    Text to token ids for a string.\n",
    "    \"\"\"\n",
    "    if type(text) is list: text = \" \".join(text)\n",
    "    output = tokenizer(text)\n",
    "    if filter_special_tokens:\n",
    "        token_ids = [i for i in output['input_ids'] if i not in tokenizer.all_special_ids]\n",
    "    else:\n",
    "        token_ids = [i for i in output['input_ids'] ]\n",
    "    return token_ids\n",
    "\n",
    "def tokenize_batch(batch):\n",
    "    \"\"\"\n",
    "    Text to token ids for a list of strings.\n",
    "    \"\"\"\n",
    "#     return [tokenizer.convert_tokens_to_ids(sent) for sent in batch]\n",
    "    return [tokenizer_return_id(sent) for sent in batch]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e58ab9de",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "get_init_text(seed_text, max_len, batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06384bfe",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "for ii in range(max_iter):\n",
    "    kk = np.random.randint(0, max_len)  # change this to id_first_mask to id_last_mask ??\n",
    "    for jj in range(batch_size):\n",
    "        batch[jj][seed_len+kk] = mask_id\n",
    "    inp = torch.tensor(batch).cuda() if cuda else torch.tensor(batch)\n",
    "    out = model(inp)\n",
    "    topk = top_k if (ii >= burnin) else 0\n",
    "    idxs = generate_step(out, gen_idx=seed_len+kk, top_k=topk, temperature=temperature, sample=(ii < burnin))\n",
    "    for jj in range(batch_size):\n",
    "        batch[jj][seed_len+kk] = idxs[jj]\n",
    "\n",
    "    if verbose and np.mod(ii+1, print_every) == 0:\n",
    "        for_print = tokenizer.convert_ids_to_tokens(batch[0])\n",
    "        for_print = for_print[:seed_len+kk+1] + ['(*)'] + for_print[seed_len+kk+1:]\n",
    "        print(\"iter\", ii+1, \" \".join(for_print))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9805dd6e",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "for ii in range(max_iter):\n",
    "    kk = np.random.randint(0, max_len)  # change this to id_first_mask to id_last_mask ??\n",
    "    for jj in range(batch_size):\n",
    "        batch[jj][seed_len+kk] = mask_id\n",
    "    inp = torch.tensor(batch).cuda() if cuda else torch.tensor(batch)\n",
    "    out = model(inp)\n",
    "    topk = top_k if (ii >= burnin) else 0\n",
    "    idxs = generate_step(out, gen_idx=seed_len+kk, top_k=topk, temperature=temperature, sample=(ii < burnin))\n",
    "    for jj in range(batch_size):\n",
    "        batch[jj][seed_len+kk] = idxs[jj]\n",
    "\n",
    "    if verbose and np.mod(ii+1, print_every) == 0:\n",
    "        for_print = tokenizer.convert_ids_to_tokens(batch[0])\n",
    "        for_print = for_print[:seed_len+kk+1] + ['(*)'] + for_print[seed_len+kk+1:]\n",
    "        print(\"iter\", ii+1, \" \".join(for_print))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5a0010d",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Original implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88d8814d",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def get_init_text(seed_text, max_len, batch_size = 1, rand_init=False):\n",
    "    \"\"\" Get initial sentence by padding seed_text with either masks or random words to max_len \"\"\"\n",
    "    batch = [seed_text + [MASK] * max_len + [SEP] for _ in range(batch_size)]\n",
    "    #if rand_init:\n",
    "    #    for ii in range(max_len):\n",
    "    #        init_idx[seed_len+ii] = np.random.randint(0, len(tokenizer.vocab))\n",
    "#     return batch\n",
    "    return tokenize_batch(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eb40e5b",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "seed_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35e43495",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "n_samples = 5\n",
    "batch_size = 5\n",
    "max_len = 40\n",
    "top_k = 100\n",
    "temperature = 1.0\n",
    "generation_mode = \"parallel-sequential\"\n",
    "leed_out_len = 5 # max_len\n",
    "burnin = 250\n",
    "sample = True\n",
    "max_iter = 500\n",
    "\n",
    "# Choose the prefix context\n",
    "seed_text = \"I have 3 more weeks for my <mask>\".split()\n",
    "seed_len = len(seed_text)\n",
    "batch = get_init_text(seed_text, 0, batch_size)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "830f088f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ecbafaf",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "for ii in range(max_len):\n",
    "    inp = [sent[:seed_len+ii+leed_out_len]+[sep_id] for sent in batch]\n",
    "    inp = torch.tensor(batch).cuda() if cuda else torch.tensor(batch)\n",
    "    out = model(inp)\n",
    "    idxs = generate_step(out, gen_idx=seed_len+ii, top_k=top_k, temperature=temperature, sample=sample)\n",
    "    for jj in range(batch_size):\n",
    "        batch[jj][seed_len+ii] = idxs[jj]\n",
    "\n",
    "return untokenize_batch(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b94bbb23",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}