preprocess_corpus/generate-TACos-multilevel-corpus.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate corpus and gruond-truth references of released videos\n",
    "\n",
    "### Corpus file contents\n",
    "0. train_data: captions and idxs of training videos in format [corpus_widxs, vidxs, corpus_pidxs], where:\n",
    "    - corpus_widxs is a list of lists with the index of words in the vocabulary\n",
    "    - vidxs is a list of indexes of video features in the features file\n",
    "    - corpus_pidxs is a list of lists with the index of POS tags in the POS tagging vocabulary\n",
    "1. val_data: same format of train_data.\n",
    "2. test_data: same format of train_data.\n",
    "3. vocabulary: in format {'word': count}.\n",
    "4. idx2word: is the vocabulary in format {idx: 'word'}.\n",
    "5. word_embeddings: are the vectors of each word. The i-th row is the word vector of the i-th word in the vocabulary.\n",
    "6. idx2pos: is the vocabulary of POS tagging in format {idx: 'POSTAG'}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate split for training and validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "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>video-id</th>\n",
       "      <th>start-frame</th>\n",
       "      <th>end-frame</th>\n",
       "      <th>sentence</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>s13-d21</td>\n",
       "      <td>185</td>\n",
       "      <td>224</td>\n",
       "      <td>the person walked into the kitchen</td>\n",
       "      <td>enter</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>s13-d21</td>\n",
       "      <td>185</td>\n",
       "      <td>224</td>\n",
       "      <td>the person entered the kitchen</td>\n",
       "      <td>enter</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>s13-d21</td>\n",
       "      <td>185</td>\n",
       "      <td>224</td>\n",
       "      <td>the person entered the kitchen</td>\n",
       "      <td>enter</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>s13-d21</td>\n",
       "      <td>185</td>\n",
       "      <td>224</td>\n",
       "      <td>the person entered a kitchen</td>\n",
       "      <td>enter</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>s13-d21</td>\n",
       "      <td>185</td>\n",
       "      <td>224</td>\n",
       "      <td>the person walked into the kitchen</td>\n",
       "      <td>enter</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>52588</th>\n",
       "      <td>s37-d74</td>\n",
       "      <td>11013</td>\n",
       "      <td>11041</td>\n",
       "      <td>the person unplugged the machine</td>\n",
       "      <td>unplug</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>52589</th>\n",
       "      <td>s37-d74</td>\n",
       "      <td>11013</td>\n",
       "      <td>11041</td>\n",
       "      <td>the person unplugged the coffee maker</td>\n",
       "      <td>unplug</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>52590</th>\n",
       "      <td>s37-d74</td>\n",
       "      <td>11013</td>\n",
       "      <td>11041</td>\n",
       "      <td>the person unplugged the machine</td>\n",
       "      <td>unplug</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>52591</th>\n",
       "      <td>s37-d74</td>\n",
       "      <td>11013</td>\n",
       "      <td>11041</td>\n",
       "      <td>the person unplugged the coffeemaker</td>\n",
       "      <td>unplug</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>52592</th>\n",
       "      <td>s37-d74</td>\n",
       "      <td>11062</td>\n",
       "      <td>11095</td>\n",
       "      <td>the person closed the outlet back up</td>\n",
       "      <td>plug in</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>52593 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      video-id  start-frame  end-frame                               sentence  \\\n",
       "0      s13-d21          185        224     the person walked into the kitchen   \n",
       "1      s13-d21          185        224         the person entered the kitchen   \n",
       "2      s13-d21          185        224         the person entered the kitchen   \n",
       "3      s13-d21          185        224           the person entered a kitchen   \n",
       "4      s13-d21          185        224     the person walked into the kitchen   \n",
       "...        ...          ...        ...                                    ...   \n",
       "52588  s37-d74        11013      11041       the person unplugged the machine   \n",
       "52589  s37-d74        11013      11041  the person unplugged the coffee maker   \n",
       "52590  s37-d74        11013      11041       the person unplugged the machine   \n",
       "52591  s37-d74        11013      11041   the person unplugged the coffeemaker   \n",
       "52592  s37-d74        11062      11095   the person closed the outlet back up   \n",
       "\n",
       "         label  \n",
       "0        enter  \n",
       "1        enter  \n",
       "2        enter  \n",
       "3        enter  \n",
       "4        enter  \n",
       "...        ...  \n",
       "52588   unplug  \n",
       "52589   unplug  \n",
       "52590   unplug  \n",
       "52591   unplug  \n",
       "52592  plug in  \n",
       "\n",
       "[52593 rows x 5 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "data = pd.read_csv('../../../data/TACoS-Multi-Level/corpus/annosDetailed-processed.csv', '\\t', usecols=[0,1,2,6,9], names=['video-id', 'start-frame', 'end-frame', 'sentence', 'label'], engine='python')  \n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# list(data['sentence'])\n",
    "\n",
    "train_vidxs, train_corpus = list(data['video-id']), list(data['sentence'])\n",
    "# valid_vidxs, valid_corpus = zip(*[(int(d['id']), d['label']) for d in valid_data])\n",
    "# test_vidxs = [(int(d['id'])) for d in test_data]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get pretrained embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "400000\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "wordvectors = {}\n",
    "# with open('./glove.42B.300d.txt') as f:\n",
    "with open('./glove.6B.300d.txt') as f:\n",
    "    for line in f:\n",
    "        s = line.strip().split(' ')\n",
    "        if len(s) == 301:\n",
    "            wordvectors[s[0]] = np.array(s[1:], dtype=float)\n",
    "    print(len(wordvectors))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Determine the vocabulary from train split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package punkt to /home/jeperez/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Avg. count of words per caption: 8.27064438233225\n",
      "Count of unique words:  2863\n",
      "missing word: refriderator\n",
      "missing word: medium-thick\n",
      "missing word: counter-top\n",
      "missing word: re-closed\n",
      "missing word: re-wrapped\n",
      "missing word: straight-edged\n",
      "missing word: cup-board\n",
      "missing word: cubbard\n",
      "missing word: dishtowel\n",
      "missing word: grapfruit\n",
      "missing word: hhe\n",
      "missing word: pieceshe\n",
      "missing word: stove-top\n",
      "missing word: drainer\n",
      "missing word: longways\n",
      "missing word: pluged\n",
      "missing word: re-sealed\n",
      "missing word: adusted\n",
      "missing word: timer/settings\n",
      "missing word: -lrb-\n",
      "missing word: -rrb-\n",
      "missing word: outter\n",
      "missing word: choped\n",
      "missing word: untoed\n",
      "missing word: un-potted\n",
      "missing word: wrraped\n",
      "missing word: diffferent\n",
      "missing word: width-wise\n",
      "missing word: length-wise\n",
      "missing word: drainboard\n",
      "missing word: potaoes\n",
      "missing word: evenly-sized\n",
      "missing word: egg-whites\n",
      "missing word: outershell\n",
      "missing word: snowpeas\n",
      "missing word: hand-washed\n",
      "missing word: de-stemmed\n",
      "missing word: flowerets\n",
      "missing word: over-cooking\n",
      "missing word: rind/eyes\n",
      "missing word: de-cored\n",
      "missing word: openned\n",
      "missing word: halfs\n",
      "missing word: infuser\n",
      "missing word: stoved\n",
      "missing word: refridgerator\n",
      "missing word: avocade\n",
      "missing word: cubing\n",
      "missing word: cubeing\n",
      "missing word: debree\n",
      "missing word: omlet\n",
      "missing word: hot-dogs\n",
      "missing word: hot-dog\n",
      "missing word: depitted\n",
      "missing word: grabed\n",
      "missing word: limejuice\n",
      "missing word: cuttingboard\n",
      "missing word: scond\n",
      "missing word: opend\n",
      "missing word: cabnent\n",
      "missing word: salt-cellar\n",
      "missing word: shrink-wrapped\n",
      "missing word: inwrapped\n",
      "missing word: broccolia\n",
      "missing word: deseed\n",
      "missing word: placd\n",
      "missing word: de-seeded\n",
      "missing word: retreaved\n",
      "missing word: trolly\n",
      "missing word: legnt-wise\n",
      "missing word: segmints\n",
      "missing word: retieed\n",
      "missing word: retrived\n",
      "missing word: bowl/strainer\n",
      "missing word: colendar\n",
      "missing word: garlics\n",
      "missing word: cuting\n",
      "missing word: avacado\n",
      "missing word: lengthways\n",
      "missing word: trhew\n",
      "missing word: rinced\n",
      "missing word: broad-bladed\n",
      "missing word: cabinet/pantry\n",
      "missing word: ginner\n",
      "missing word: herinsed\n",
      "missing word: chip-sized\n",
      "missing word: warapper\n",
      "missing word: brocoli\n",
      "missing word: flowerettes\n",
      "missing word: inch-long\n",
      "missing word: baord\n",
      "missing word: 2cm\n",
      "missing word: cuttting\n",
      "missing word: maqngo\n",
      "missing word: overwrap\n",
      "missing word: rinded\n",
      "missing word: rewashed\n",
      "missing word: quatered\n",
      "missing word: egss\n",
      "missing word: retreaveeed\n",
      "missing word: cubboard\n",
      "missing word: frigerator\n",
      "missing word: laddle\n",
      "missing word: ingreadiants\n",
      "missing word: a455\n",
      "missing word: stired\n",
      "missing word: ontop\n",
      "missing word: hand-juicer\n",
      "missing word: plup\n",
      "missing word: a439\n",
      "missing word: weiners\n",
      "missing word: cucumbe\n",
      "missing word: breifly\n",
      "missing word: tooked\n",
      "missing word: choppinned\n",
      "missing word: woodedn\n",
      "missing word: etiolated\n",
      "missing word: plced\n",
      "missing word: vegi\n",
      "missing word: asmall\n",
      "missing word: peices\n",
      "missing word: medium-width\n",
      "missing word: middled\n",
      "missing word: bited\n",
      "missing word: half-loaf\n",
      "missing word: re-twisted\n",
      "missing word: toastd\n",
      "missing word: amnd\n",
      "missing word: anothe\n",
      "missing word: buter\n",
      "missing word: ingridients\n",
      "missing word: bottles/packetted\n",
      "missing word: 7minutes\n",
      "missing word: contiuned\n",
      "missing word: cupard\n",
      "missing word: comapartment\n",
      "missing word: slected\n",
      "missing word: cubbord\n",
      "missing word: legnth\n",
      "missing word: makeeed\n",
      "missing word: derection\n",
      "missing word: brocali\n",
      "missing word: board/knife\n",
      "missing word: useeed\n",
      "missing word: inchs\n",
      "missing word: broccolid\n",
      "missing word: brocalli\n",
      "missing word: ocasionally\n",
      "missing word: eventaully\n",
      "missing word: strains/rinsed\n",
      "missing word: placeeed\n",
      "missing word: saltshaker\n",
      "missing word: egg-white\n",
      "missing word: egg-yolk\n",
      "missing word: stem-like\n",
      "missing word: opned\n",
      "missing word: slized\n",
      "missing word: slizes\n",
      "missing word: lengths-wise\n",
      "missing word: woaman\n",
      "missing word: one-fourths\n",
      "missing word: cutting-board\n",
      "missing word: left-overs\n",
      "missing word: arils\n",
      "missing word: peel/skin\n",
      "missing word: skined\n",
      "missing word: aparatus\n",
      "missing word: pan-frying\n",
      "missing word: fliped\n",
      "missing word: piecses\n",
      "missing word: sauted\n",
      "missing word: shre\n",
      "missing word: microplaner\n",
      "missing word: refridgeator\n",
      "missing word: quaters\n",
      "missing word: kichen\n",
      "missing word: skin-less\n",
      "missing word: water-kettle\n",
      "missing word: occaionally\n",
      "missing word: cubard\n",
      "missing word: refrigertor\n",
      "missing word: rubber-tipped\n",
      "missing word: towel-dried\n",
      "missing word: verticle\n",
      "missing word: piecess\n",
      "missing word: downwarded\n",
      "missing word: diced-size\n",
      "missing word: claw-shaped\n",
      "missing word: taste-tested\n",
      "missing word: a283\n",
      "missing word: eggcup\n",
      "missing word: pealer\n",
      "missing word: extracter\n",
      "missing word: citrus-juicer\n",
      "missing word: coverd\n",
      "missing word: toater\n",
      "missing word: two-slice\n",
      "missing word: regirgerator\n",
      "missing word: dryed\n",
      "missing word: unpluged\n",
      "missing word: katchip\n",
      "missing word: stuck-on\n",
      "missing word: tiffun\n",
      "missing word: juiceless\n",
      "missing word: pomegeranate\n",
      "missing word: deseeded\n",
      "missing word: plastic-wrapped\n",
      "missing word: now-chopped\n",
      "missing word: wealked\n",
      "missing word: hotplate\n",
      "missing word: removeeed\n",
      "missing word: rinseeed\n",
      "missing word: earlyer\n",
      "missing word: cloved\n",
      "missing word: siled\n",
      "missing word: enterded\n",
      "missing word: drawerred\n",
      "missing word: chakla\n",
      "missing word: dhaniya\n",
      "missing word: orignal\n",
      "missing word: re-approached\n",
      "missing word: re-rinsed\n",
      "missing word: re-washed\n",
      "missing word: toaste\n",
      "missing word: knif\n",
      "missing word: longwise\n",
      "missing word: one-centimeter\n",
      "missing word: straighted\n",
      "missing word: retreived\n",
      "missing word: leeked\n",
      "missing word: under-counter\n",
      "missing word: drip-style\n",
      "missing word: un-bleached\n",
      "count of missing words:  233\n",
      "2630 2632 2632\n"
     ]
    }
   ],
   "source": [
    "import nltk\n",
    "nltk.download('punkt')\n",
    "\n",
    "vocab, total_len = {}, 0\n",
    "for cap in train_corpus:\n",
    "    tokens = nltk.word_tokenize(cap.lower())\n",
    "    total_len += len(tokens)\n",
    "    for w in tokens:\n",
    "        try:\n",
    "            vocab[w] += 1\n",
    "        except:\n",
    "            vocab[w] = 1\n",
    "\n",
    "print('Avg. count of words per caption:', total_len/len(train_corpus))\n",
    "print('Count of unique words: ', len(vocab))\n",
    "\n",
    "to_del = []\n",
    "for w in vocab.keys():\n",
    "    if not w in wordvectors:\n",
    "        to_del.append(w)\n",
    "        print('missing word: {}'.format(w))\n",
    "\n",
    "print('count of missing words: ', len(to_del))\n",
    "        \n",
    "for w in to_del:\n",
    "    del vocab[w]\n",
    "        \n",
    "idx2word = {idx: word for idx, word in enumerate(['<eos>', '<unk>'] + list(vocab.keys()))}\n",
    "word2idx = {word: idx for idx, word in enumerate(['<eos>', '<unk>'] + list(vocab.keys()))}\n",
    "EOS, UNK = 0, 1\n",
    "\n",
    "print(len(vocab), len(idx2word), len(word2idx))\n",
    "\n",
    "word_embeddings = np.zeros((len(idx2word), 300))\n",
    "for idx, word in idx2word.items():\n",
    "    if idx == EOS:\n",
    "        word_embeddings[idx] = wordvectors['eos']\n",
    "    elif idx == UNK:\n",
    "        word_embeddings[idx] = wordvectors['unk']\n",
    "    else:\n",
    "        word_embeddings[idx] = wordvectors[word]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Determine POS-tagging vocabulary from train split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unique words per tag:\n",
      " DT:\t19\n",
      " NN:\t1178\n",
      " VBD:\t577\n",
      " IN:\t76\n",
      " NNS:\t347\n",
      " VBG:\t198\n",
      " CC:\t6\n",
      " JJ:\t582\n",
      " TO:\t1\n",
      " VB:\t306\n",
      " RP:\t15\n",
      " ,:\t1\n",
      " JJR:\t28\n",
      " RB:\t196\n",
      " PRP$:\t5\n",
      " VBN:\t172\n",
      " CD:\t22\n",
      " PRP:\t11\n",
      " FW:\t6\n",
      " VBP:\t65\n",
      " WDT:\t3\n",
      " RBR:\t11\n",
      " VBZ:\t80\n",
      " PDT:\t6\n",
      " POS:\t2\n",
      " MD:\t5\n",
      " JJS:\t9\n",
      " EX:\t1\n",
      " WRB:\t2\n",
      " WP:\t3\n",
      " NNP:\t6\n",
      " ::\t3\n",
      " ``:\t2\n",
      " .:\t2\n",
      " RBS:\t1\n",
      " '':\t1\n",
      "38\n"
     ]
    }
   ],
   "source": [
    "import nltk\n",
    "\n",
    "pos_vocab = {}\n",
    "pos_unique_words = {}\n",
    "for cap in train_corpus:\n",
    "    for tag in nltk.pos_tag(nltk.word_tokenize(cap.lower())):\n",
    "        try:\n",
    "            pos_vocab[tag[1]] += 1\n",
    "            try: \n",
    "                pos_unique_words[tag[1]][tag[0]] += 1\n",
    "            except:\n",
    "                pos_unique_words[tag[1]][tag[0]] = 1\n",
    "        except:\n",
    "            pos_vocab[tag[1]] = 1\n",
    "            pos_unique_words[tag[1]] = {tag[0]: 1}\n",
    "\n",
    "print('Unique words per tag:')\n",
    "print('\\n'.join([f' {k}:\\t{len(words)}' for k, words in pos_unique_words.items()]))\n",
    "            \n",
    "idx2pos = {idx: tag for idx, tag in enumerate(['eos', 'unk'] + list(pos_vocab.keys()))}\n",
    "pos2idx = {tag: idx for idx, tag in enumerate(['eos', 'unk'] + list(pos_vocab.keys()))}\n",
    "EOS, UNK = 0, 1\n",
    "print(len(idx2pos))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Determine Universal POS-tagging from train split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package universal_tagset to\n",
      "[nltk_data]     /home/jeperez/nltk_data...\n",
      "[nltk_data]   Package universal_tagset is already up-to-date!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unique words per universal tag:\n",
      " DET:\t26\n",
      " NOUN:\t1475\n",
      " VERB:\t1178\n",
      " ADP:\t76\n",
      " CONJ:\t6\n",
      " ADJ:\t609\n",
      " PRT:\t18\n",
      " .:\t9\n",
      " ADV:\t207\n",
      " PRON:\t19\n",
      " NUM:\t22\n",
      " X:\t6\n",
      "14\n"
     ]
    }
   ],
   "source": [
    "import nltk\n",
    "nltk.download('universal_tagset')\n",
    "\n",
    "upos_vocab = {}\n",
    "upos_unique_words = {}\n",
    "for cap in train_corpus:\n",
    "    for tag in nltk.pos_tag(nltk.word_tokenize(cap.lower()), tagset='universal'):\n",
    "        try:\n",
    "            upos_vocab[tag[1]] += 1\n",
    "            try: \n",
    "                upos_unique_words[tag[1]][tag[0]] += 1\n",
    "            except:\n",
    "                upos_unique_words[tag[1]][tag[0]] = 1\n",
    "        except:\n",
    "            upos_vocab[tag[1]] = 1\n",
    "            upos_unique_words[tag[1]] = {tag[0]: 1}\n",
    "\n",
    "print('Unique words per universal tag:')\n",
    "print('\\n'.join([f' {k}:\\t{len(words)}' for k, words in upos_unique_words.items()]))\n",
    "            \n",
    "idx2upos = {idx: word for idx, word in enumerate(['eos', 'unk'] + list(upos_vocab.keys()))}\n",
    "upos2idx = {word: idx for idx, word in enumerate(['eos', 'unk'] + list(upos_vocab.keys()))}\n",
    "print(len(idx2upos))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate ground-truth references files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../results/20B-SS-v2_val_references.txt', 'w') as f:\n",
    "    for vidx, cap in zip(valid_vidxs, valid_corpus):\n",
    "        f.write('{}\\t{}\\n'.format(vidx, cap.lower()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate corpus.pkl file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "train_corpus_widxs = [[word2idx[w] if w in vocab else UNK for w in nltk.word_tokenize(cap.lower())] + [EOS] for cap in train_corpus]\n",
    "valid_corpus_widxs = [[word2idx[w] if w in vocab else UNK for w in nltk.word_tokenize(cap.lower())] + [EOS] for cap in valid_corpus]\n",
    "\n",
    "train_corpus_pidxs = [[pos2idx[w[1]] if w[1] in pos_vocab else UNK for w in nltk.pos_tag(nltk.word_tokenize(cap.lower()))] + [EOS] for cap in train_corpus]\n",
    "valid_corpus_pidxs = [[pos2idx[w[1]] if w[1] in pos_vocab else UNK for w in nltk.pos_tag(nltk.word_tokenize(cap.lower()))] + [EOS] for cap in valid_corpus]\n",
    "\n",
    "assert len(train_corpus_widxs) == len(train_vidxs) and len(train_vidxs) == len(train_corpus_pidxs) and len(train_vidxs) == len(train_corpus)\n",
    "assert len(valid_corpus_widxs) == len(valid_vidxs) and len(valid_vidxs) == len(valid_corpus_pidxs) and len(valid_vidxs) == len(valid_corpus)\n",
    "\n",
    "train_data = [train_corpus_widxs, train_vidxs, train_corpus_pidxs, train_corpus]\n",
    "valid_data = [valid_corpus_widxs, valid_vidxs, valid_corpus_pidxs, valid_corpus]\n",
    "test_data = [None, test_vidxs, None]\n",
    "\n",
    "with open('../../../data/Something-Something-v2/20b-ss-v2_corpus_pos.pkl', 'wb') as outfile:\n",
    "    pickle.dump([train_data, valid_data, test_data, vocab, idx2word, word_embeddings, idx2pos], outfile)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}