Beam Decoding.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Beam decoding is essential for seq2seq tasks. But it's notoriously complicated to implement. Here's a relatively easy one, batchfying candidates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "__author__ = \"kyubyong\"\n",
    "__address__ = \"https://github.com/kyubyong/nlp_made_easy\"\n",
    "__email__ = \"kbpark.linguist@gmail.com\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 262,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = 100 # output dimensionality. number of vocabulary\n",
    "H = 10 # hidden dimensionality \n",
    "K = 3 # beam width\n",
    "T = 10 # decoding timesteps\n",
    "initial_y = np.array([[3], [6]], dtype=np.int32)\n",
    "N = len(initial_y) # batch size\n",
    "xh = np.random.randn(V, H) # weights from input to hidden\n",
    "hh = np.random.randn(H, H) # weights from hidden to hidden\n",
    "ho = np.random.randn(H, V) # weights from hidden to outputs\n",
    "\n",
    "EOS_ID = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 263,
   "metadata": {},
   "outputs": [],
   "source": [
    "def onehot(arry, size):\n",
    "    '''\n",
    "    arry: 2-d array of n, t\n",
    "    size: output dimensions\n",
    "    \n",
    "    returns\n",
    "    3-d array of (n, t, size)\n",
    "    '''\n",
    "    labels_one_hot = (arry.ravel()[np.newaxis] == np.arange(size)[:, np.newaxis]).T\n",
    "    labels_one_hot.shape = arry.shape + (size,)\n",
    "    return labels_one_hot.astype('int32')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 264,
   "metadata": {},
   "outputs": [],
   "source": [
    "def softmax(x):\n",
    "    \"\"\"Compute softmax values for each sets of scores in x.\"\"\"\n",
    "    e_x = np.exp(x - np.max(x))\n",
    "    return e_x / e_x.sum(axis=1,keepdims=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 270,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decode(y):\n",
    "    '''Decodes with a simple rnn.\n",
    "    y: (N, t) array if not None.\n",
    "    '''\n",
    "    global N, V, H, xh, hh, ho\n",
    "    \n",
    "    prev_hidden = np.zeros((y.shape[0], H)) # initial hidden\n",
    "    \n",
    "    for t in range(y.shape[1]):\n",
    "        token = y[:, t]\n",
    "        x_to_h = np.matmul(onehot(token, V), xh) # (N, h)\n",
    "        h_to_h = np.matmul(prev_hidden, hh) # (N, h)\n",
    "        hidden = np.tanh(x_to_h + h_to_h)\n",
    "        prev_hidden = hidden\n",
    "        if t == y.shape[1]-1: # last step\n",
    "            outputs = np.matmul(hidden, ho)\n",
    "            probs = softmax(outputs)\n",
    "            probs = np.log(probs)\n",
    "            return probs # (N, V)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Beam decode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 268,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "========== timesteps= 0 ==========\n",
      "batch num= 0\n",
      "[13]\tFalse\t-1.801408656139552\n",
      "[48]\tFalse\t-1.9670540723462346\n",
      "[33]\tFalse\t-2.198485320487616\n",
      "batch num= 1\n",
      "[30]\tFalse\t-0.7897988012958753\n",
      "[1]\tFalse\t-2.4629262078330942\n",
      "[76]\tFalse\t-2.942492396759654\n",
      "========== timesteps= 1 ==========\n",
      "Expansion...\n",
      "batch num= 0\n",
      "[13 53]\tFalse\t-1.7223380770456393\n",
      "[13 13]\tFalse\t-2.153169837293493\n",
      "[13 21]\tFalse\t-2.1963759327725914\n",
      "[48  7]\tFalse\t-1.5501885837535672\n",
      "[48 56]\tFalse\t-2.0774770760291426\n",
      "[48 71]\tFalse\t-2.2559148677511205\n",
      "[33 13]\tFalse\t-2.02571782070014\n",
      "[33 20]\tFalse\t-2.2622432967465382\n",
      "[33 99]\tFalse\t-2.411861897194325\n",
      "batch num= 1\n",
      "[30  8]\tFalse\t-1.0387280196474995\n",
      "[30 38]\tFalse\t-1.4677326159171868\n",
      "[30 67]\tFalse\t-1.963540864046755\n",
      "[ 1 46]\tFalse\t-1.8347386403083892\n",
      "[ 1 50]\tFalse\t-2.062073079883135\n",
      "[ 1 71]\tFalse\t-2.79060930087043\n",
      "[76 89]\tFalse\t-2.4038844011808673\n",
      "[76 44]\tFalse\t-2.440214472610669\n",
      "[76  3]\tFalse\t-2.4820893854715784\n",
      "Pruning ...\n",
      "batch num= 0\n",
      "[48  7]\tFalse\t-1.5501885837535672\n",
      "[13 53]\tFalse\t-1.7223380770456393\n",
      "[33 13]\tFalse\t-2.02571782070014\n",
      "batch num= 1\n",
      "[30  8]\tFalse\t-1.0387280196474995\n",
      "[30 38]\tFalse\t-1.4677326159171868\n",
      "[ 1 46]\tFalse\t-1.8347386403083892\n",
      "========== timesteps= 2 ==========\n",
      "Expansion...\n",
      "batch num= 0\n",
      "[48  7 45]\tFalse\t-1.4907133631014375\n",
      "[48  7 49]\tFalse\t-1.5125758679961583\n",
      "[48  7 46]\tFalse\t-1.68788305167278\n",
      "[13 53 53]\tFalse\t-1.4720320535925133\n",
      "[13 53 17]\tFalse\t-1.503669670852281\n",
      "[13 53 49]\tFalse\t-2.2055456936009072\n",
      "[33 13 53]\tFalse\t-1.868266739269855\n",
      "[33 13 17]\tFalse\t-1.9480947954999472\n",
      "[33 13 16]\tFalse\t-2.0038997227337942\n",
      "batch num= 1\n",
      "[30  8  7]\tFalse\t-0.981799161726694\n",
      "[30  8 50]\tFalse\t-1.2210058993765245\n",
      "[30  8 33]\tFalse\t-1.6862021730157182\n",
      "[30 38  7]\tFalse\t-1.0155674726535289\n",
      "[30 38 50]\tFalse\t-2.4361530449230084\n",
      "[30 38 56]\tFalse\t-2.4468657026464666\n",
      "[ 1 46 43]\tFalse\t-1.730098315254036\n",
      "[ 1 46 23]\tFalse\t-1.910274720612243\n",
      "[ 1 46 64]\tFalse\t-1.9826010216276544\n",
      "Pruning ...\n",
      "batch num= 0\n",
      "[13 53 53]\tFalse\t-1.4720320535925133\n",
      "[48  7 45]\tFalse\t-1.4907133631014375\n",
      "[13 53 17]\tFalse\t-1.503669670852281\n",
      "batch num= 1\n",
      "[30  8  7]\tFalse\t-0.981799161726694\n",
      "[30 38  7]\tFalse\t-1.0155674726535289\n",
      "[30  8 50]\tFalse\t-1.2210058993765245\n",
      "========== timesteps= 3 ==========\n",
      "Expansion...\n",
      "batch num= 0\n",
      "[13 53 53 41]\tFalse\t-1.5506704012318269\n",
      "[13 53 53 25]\tFalse\t-1.561838485804779\n",
      "[13 53 53 98]\tFalse\t-1.583598742175906\n",
      "[48  7 45 80]\tFalse\t-1.435019828867431\n",
      "[48  7 45 76]\tFalse\t-1.44588109619326\n",
      "[48  7 45  1]\tFalse\t-1.7402676634547847\n",
      "[13 53 17 41]\tFalse\t-1.585756130334546\n",
      "[13 53 17 80]\tFalse\t-1.5910758874599538\n",
      "[13 53 17 25]\tFalse\t-1.6346056800788564\n",
      "batch num= 1\n",
      "[30  8  7 43]\tFalse\t-1.0350845938886057\n",
      "[30  8  7 69]\tFalse\t-1.487864532678162\n",
      "[30  8  7 14]\tFalse\t-1.4976892653945955\n",
      "[30 38  7 46]\tFalse\t-0.8816524694126934\n",
      "[30 38  7 45]\tFalse\t-1.361845258654082\n",
      "[30 38  7 50]\tFalse\t-1.3850981222623515\n",
      "[30  8 50 43]\tFalse\t-1.474636461781719\n",
      "[30  8 50 11]\tFalse\t-1.6120821984876703\n",
      "[30  8 50 14]\tFalse\t-1.6380255433340407\n",
      "Pruning ...\n",
      "batch num= 0\n",
      "[48  7 45 80]\tFalse\t-1.435019828867431\n",
      "[48  7 45 76]\tFalse\t-1.44588109619326\n",
      "[13 53 53 41]\tFalse\t-1.5506704012318269\n",
      "batch num= 1\n",
      "[30 38  7 46]\tFalse\t-0.8816524694126934\n",
      "[30  8  7 43]\tFalse\t-1.0350845938886057\n",
      "[30 38  7 45]\tFalse\t-1.361845258654082\n",
      "========== timesteps= 4 ==========\n",
      "Expansion...\n",
      "batch num= 0\n",
      "[48  7 45 80  1]\tFalse\t-1.4155848839072929\n",
      "[48  7 45 80 63]\tFalse\t-1.4809062680622915\n",
      "[48  7 45 80 94]\tFalse\t-1.732865605759868\n",
      "[48  7 45 76 67]\tFalse\t-1.581737929223578\n",
      "[48  7 45 76 92]\tFalse\t-1.5835056363202569\n",
      "[48  7 45 76 13]\tFalse\t-1.7110238710735544\n",
      "[13 53 53 41 62]\tFalse\t-1.417947578093456\n",
      "[13 53 53 41 66]\tFalse\t-1.5861790359814807\n",
      "[13 53 53 41 22]\tFalse\t-1.8198367051022502\n",
      "batch num= 1\n",
      "[30 38  7 46 15]\tFalse\t-0.8900240329071243\n",
      "[30 38  7 46 76]\tFalse\t-1.0029480860843567\n",
      "[30 38  7 46 40]\tFalse\t-1.2061234517379482\n",
      "[30  8  7 43 40]\tFalse\t-1.1621896355270631\n",
      "[30  8  7 43 74]\tFalse\t-1.1756940719777458\n",
      "[30  8  7 43 76]\tFalse\t-1.2723109905838466\n",
      "[30 38  7 45 76]\tFalse\t-1.2788018546331241\n",
      "[30 38  7 45 15]\tFalse\t-1.4225672540053953\n",
      "[30 38  7 45 40]\tFalse\t-1.5114030374152683\n",
      "Pruning ...\n",
      "batch num= 0\n",
      "[48  7 45 80  1]\tFalse\t-1.4155848839072929\n",
      "[13 53 53 41 62]\tFalse\t-1.417947578093456\n",
      "[48  7 45 80 63]\tFalse\t-1.4809062680622915\n",
      "batch num= 1\n",
      "[30 38  7 46 15]\tFalse\t-0.8900240329071243\n",
      "[30 38  7 46 76]\tFalse\t-1.0029480860843567\n",
      "[30  8  7 43 40]\tFalse\t-1.1621896355270631\n",
      "========== timesteps= 5 ==========\n",
      "Expansion...\n",
      "batch num= 0\n",
      "[48  7 45 80  1  7]\tFalse\t-1.4794142075137813\n",
      "[48  7 45 80  1 41]\tFalse\t-1.4933938451514475\n",
      "[48  7 45 80  1 71]\tFalse\t-1.5103248587788836\n",
      "[13 53 53 41 62 76]\tFalse\t-1.3396238306377644\n",
      "[13 53 53 41 62 89]\tFalse\t-1.4990384205012208\n",
      "[13 53 53 41 62 65]\tFalse\t-1.5266813295887929\n",
      "[48  7 45 80 63 38]\tFalse\t-1.5578383956446686\n",
      "[48  7 45 80 63 71]\tFalse\t-1.5693606690867572\n",
      "[48  7 45 80 63 67]\tFalse\t-1.5881997528395313\n",
      "batch num= 1\n",
      "[30 38  7 46 15  1]\tFalse\t-1.0160237052072278\n",
      "[30 38  7 46 15  0]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15 95]\tFalse\t-1.1659301187826023\n",
      "[30 38  7 46 76 79]\tFalse\t-1.2324967288788231\n",
      "[30 38  7 46 76 92]\tFalse\t-1.242909393188552\n",
      "[30 38  7 46 76 75]\tFalse\t-1.3027834896100112\n",
      "[30  8  7 43 40 43]\tFalse\t-1.2051650221374317\n",
      "[30  8  7 43 40 14]\tFalse\t-1.2439917222130705\n",
      "[30  8  7 43 40 22]\tFalse\t-1.323827336732793\n",
      "Pruning ...\n",
      "batch num= 0\n",
      "[13 53 53 41 62 76]\tFalse\t-1.3396238306377644\n",
      "[48  7 45 80  1  7]\tFalse\t-1.4794142075137813\n",
      "[48  7 45 80  1 41]\tFalse\t-1.4933938451514475\n",
      "batch num= 1\n",
      "[30 38  7 46 15  0]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  1]\tFalse\t-1.0160237052072278\n",
      "[30 38  7 46 15 95]\tFalse\t-1.1659301187826023\n",
      "========== timesteps= 6 ==========\n",
      "Expansion...\n",
      "batch num= 0\n",
      "[13 53 53 41 62 76 39]\tFalse\t-1.3496708801476172\n",
      "[13 53 53 41 62 76  8]\tFalse\t-1.4023681944692268\n",
      "[13 53 53 41 62 76 23]\tFalse\t-1.4492251873156643\n",
      "[48  7 45 80  1  7 82]\tFalse\t-1.5279394757013531\n",
      "[48  7 45 80  1  7 45]\tFalse\t-1.5425022838809215\n",
      "[48  7 45 80  1  7 46]\tFalse\t-1.5619046177353486\n",
      "[48  7 45 80  1 41 82]\tFalse\t-1.430451510587661\n",
      "[48  7 45 80  1 41 68]\tFalse\t-1.6258046549927043\n",
      "[48  7 45 80  1 41 23]\tFalse\t-1.6287973043777444\n",
      "batch num= 1\n",
      "[30 38  7 46 15  0 67]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 20]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 57]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  1 11]\tFalse\t-1.0339379651820644\n",
      "[30 38  7 46 15  1 76]\tFalse\t-1.1711319401110587\n",
      "[30 38  7 46 15  1 41]\tFalse\t-1.2156144563727591\n",
      "[30 38  7 46 15 95 20]\tFalse\t-1.2206017964230695\n",
      "[30 38  7 46 15 95 11]\tFalse\t-1.2973599838246521\n",
      "[30 38  7 46 15 95 39]\tFalse\t-1.347143453699934\n",
      "Pruning ...\n",
      "batch num= 0\n",
      "[13 53 53 41 62 76 39]\tFalse\t-1.3496708801476172\n",
      "[13 53 53 41 62 76  8]\tFalse\t-1.4023681944692268\n",
      "[48  7 45 80  1 41 82]\tFalse\t-1.430451510587661\n",
      "batch num= 1\n",
      "[30 38  7 46 15  0 57]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 20]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67]\tTrue\t-0.8900240329071243\n",
      "========== timesteps= 7 ==========\n",
      "Expansion...\n",
      "batch num= 0\n",
      "[13 53 53 41 62 76 39 78]\tFalse\t-1.4009193941899798\n",
      "[13 53 53 41 62 76 39 38]\tFalse\t-1.4184446787212193\n",
      "[13 53 53 41 62 76 39 72]\tFalse\t-1.4859004214067744\n",
      "[13 53 53 41 62 76  8 20]\tFalse\t-1.4317120180607932\n",
      "[13 53 53 41 62 76  8 79]\tFalse\t-1.4555651089534953\n",
      "[13 53 53 41 62 76  8 88]\tFalse\t-1.5136947745882172\n",
      "[48  7 45 80  1 41 82 15]\tFalse\t-1.4773084124370452\n",
      "[48  7 45 80  1 41 82 84]\tFalse\t-1.5052793912158204\n",
      "[48  7 45 80  1 41 82 30]\tFalse\t-1.5128248163630005\n",
      "batch num= 1\n",
      "[30 38  7 46 15  0 57  8]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 57 22]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 57 23]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 20  8]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 20 22]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 20 62]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 56]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67  8]\tTrue\t-0.8900240329071243\n",
      "Pruning ...\n",
      "batch num= 0\n",
      "[13 53 53 41 62 76 39 78]\tFalse\t-1.4009193941899798\n",
      "[13 53 53 41 62 76 39 38]\tFalse\t-1.4184446787212193\n",
      "[13 53 53 41 62 76  8 20]\tFalse\t-1.4317120180607932\n",
      "batch num= 1\n",
      "[30 38  7 46 15  0 67  8]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 56]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82]\tTrue\t-0.8900240329071243\n",
      "========== timesteps= 8 ==========\n",
      "Expansion...\n",
      "batch num= 0\n",
      "[13 53 53 41 62 76 39 78  7]\tFalse\t-1.4068256900320346\n",
      "[13 53 53 41 62 76 39 78 49]\tFalse\t-1.4296420137168409\n",
      "[13 53 53 41 62 76 39 78 45]\tFalse\t-1.4415205750559386\n",
      "[13 53 53 41 62 76 39 38  7]\tFalse\t-1.3760183790382272\n",
      "[13 53 53 41 62 76 39 38 45]\tFalse\t-1.5019493281169982\n",
      "[13 53 53 41 62 76 39 38 50]\tFalse\t-1.511051096358533\n",
      "[13 53 53 41 62 76  8 20 79]\tFalse\t-1.4413910156791327\n",
      "[13 53 53 41 62 76  8 20 62]\tFalse\t-1.5197912451454039\n",
      "[13 53 53 41 62 76  8 20 66]\tFalse\t-1.522338865864498\n",
      "batch num= 1\n",
      "[30 38  7 46 15  0 67  8 76]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67  8 15]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67  8 40]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 56 56]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 56 87]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 56 69]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82  7]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82 56]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82 82]\tTrue\t-0.8900240329071243\n",
      "Pruning ...\n",
      "batch num= 0\n",
      "[13 53 53 41 62 76 39 38  7]\tFalse\t-1.3760183790382272\n",
      "[13 53 53 41 62 76 39 78  7]\tFalse\t-1.4068256900320346\n",
      "[13 53 53 41 62 76 39 78 49]\tFalse\t-1.4296420137168409\n",
      "batch num= 1\n",
      "[30 38  7 46 15  0 67 82 82]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82 56]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82  7]\tTrue\t-0.8900240329071243\n",
      "========== timesteps= 9 ==========\n",
      "Expansion...\n",
      "batch num= 0\n",
      "[13 53 53 41 62 76 39 38  7 46]\tFalse\t-1.2724431112719832\n",
      "[13 53 53 41 62 76 39 38  7 50]\tFalse\t-1.472136698673575\n",
      "[13 53 53 41 62 76 39 38  7 45]\tFalse\t-1.5468339153451782\n",
      "[13 53 53 41 62 76 39 78  7 46]\tFalse\t-1.3125356502306427\n",
      "[13 53 53 41 62 76 39 78  7 45]\tFalse\t-1.513518005083196\n",
      "[13 53 53 41 62 76 39 78  7 49]\tFalse\t-1.5177294947988433\n",
      "[13 53 53 41 62 76 39 78 49 49]\tFalse\t-1.4052002087742437\n",
      "[13 53 53 41 62 76 39 78 49 45]\tFalse\t-1.4535441308184314\n",
      "[13 53 53 41 62 76 39 78 49 46]\tFalse\t-1.5499813329569079\n",
      "batch num= 1\n",
      "[30 38  7 46 15  0 67 82 82 49]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82 82 15]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82 82 45]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82 56 15]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82 56 46]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82 56 49]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82  7 46]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82  7 45]\tTrue\t-0.8900240329071243\n",
      "[30 38  7 46 15  0 67 82  7 49]\tTrue\t-0.8900240329071243\n",
      "Pruning ...\n",
      "batch num= 0\n",
      "[13 53 53 41 62 76 39 38  7 46]\tFalse\t-1.2724431112719832\n",
      "batch num= 1\n",
      "[30 38  7 46 15  0 67 82  7 49]\tTrue\t-0.8900240329071243\n"
     ]
    }
   ],
   "source": [
    "for t in range(T):  \n",
    "    def _get_preds_and_probs(PREDS):\n",
    "        probs = decode(y=PREDS) # (N, V)\n",
    "        preds_k = np.argsort(probs)[:, ::-1][:, :K].flatten() # (K*N,)\n",
    "        probs_k = np.sort(probs)[:, ::-1][:, :K].flatten() # (K*N,)\n",
    "        return preds_k, probs_k\n",
    "    \n",
    "    def logging(PREDS_k, EOS_k, PROBS_k):\n",
    "        for i, (PREDS_k_batch, EOS_k_batch, PROBS_k_batch) in \\\n",
    "                    enumerate(zip(np.split(PREDS_k, N), np.split(EOS_k, N), np.split(PROBS_k, N) )):\n",
    "            print(\"batch num=\", i)\n",
    "            for each_PREDS_k_batch, each_EOS_k_batch, each_PROBS_k_batch in zip(PREDS_k_batch, EOS_k_batch, PROBS_k_batch):\n",
    "                print(\"{}\\t{}\\t{}\".format(each_PREDS_k_batch, each_EOS_k_batch, each_PROBS_k_batch))\n",
    "               \n",
    "    if t==0: # initial step\n",
    "        print(\"=\"*10, \"timesteps=\", t, \"=\"*10)\n",
    "        \n",
    "        preds_k, probs_k = _get_preds_and_probs(initial_y) # (k*N), (k*N)\n",
    "        PREDS_k = np.expand_dims(preds_k, -1) # PREDS_k: Final outputs, (k*N, 1)\n",
    "        PROBS_k = probs_k\n",
    "        EOS_k = preds_k==EOS_ID # \n",
    "        \n",
    "        # logging\n",
    "        logging(PREDS_k, EOS_k, PROBS_k)\n",
    "                                                      \n",
    "    else:\n",
    "        print(\"=\"*10, \"timesteps=\", t, \"=\"*10)\n",
    "        print(\"Expansion...\")\n",
    "        \n",
    "        preds_kk, probs_kk = _get_preds_and_probs(PREDS_k) # (k*k*N), (k*k*N) <- incremental(=local) values\n",
    "        \n",
    "        # preds for exanded beams\n",
    "        PREDS_kk = np.repeat(PREDS_k, k, axis=0) # (k*k*N, t)\n",
    "        PREDS_kk = np.append(PREDS_kk, np.expand_dims(preds_kk, -1), -1) # PREDS_kk: (k*k*N, t+1)\n",
    "        \n",
    "        # eos for expanded beams\n",
    "        eos_kk = preds_kk==EOS_ID # (k*k*N) <- local\n",
    "        EOS_kk = np.repeat(EOS_k, k, axis=0) # (k*k*N, )\n",
    "        EOS_kk = np.logical_or(EOS_kk, eos_kk) # (k*k*N,)\n",
    "        \n",
    "        # probs for expanded beams\n",
    "        PROBS_kk = np.repeat(PROBS_k, k, axis=0) # (k*k*N, )\n",
    "        normalized_probs = ( PROBS_kk*t + probs_kk ) / (t+1)\n",
    "        PROBS_kk = np.where(EOS_kk, PROBS_kk, normalized_probs) # (k*k*N, )\n",
    "        \n",
    "        # logging\n",
    "        logging(PREDS_kk, EOS_kk, PROBS_kk)\n",
    "        \n",
    "        print(\"Pruning ...\")\n",
    "        winners = [] # (k*N). k elements are selected out of k^2\n",
    "        for j, prob_kk in enumerate(np.split(PROBS_kk, N)): # (k*k,) \n",
    "            if t == T-1: # final step\n",
    "                winner = np.argsort(prob_kk)[::-1][:1] # final 1 best\n",
    "                winners.extend(list(winner + j*len(prob_kk)))\n",
    "            else:\n",
    "                winner = np.argsort(prob_kk)[::-1][:k]\n",
    "                winners.extend(list(winner + j*len(prob_kk)))\n",
    "        \n",
    "        PREDS_k = PREDS_kk[winners] # (N, T) if final step,  otherwise (k*N, t)\n",
    "        PROBS_k = PROBS_kk[winners] # (N, T) if final step,  otherwise (k*N, )\n",
    "        EOS_k = EOS_kk[winners]\n",
    "        \n",
    "        # logging\n",
    "        logging(PREDS_k, EOS_k, PROBS_k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Be aware that the tokens that follows the `<EOS>` are stripped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}