inferences.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "inferences.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyN+iuw4rL7Tj2kkhn2HDlIq",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/DmitriyValetov/nlp_course_project/blob/master/inferences.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D4r5ZXvKsFx4",
        "colab_type": "code",
        "outputId": "8f0a0887-f10d-4204-9b51-5acc5f2217cd",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 232
        }
      },
      "source": [
        "import numpy as np\n",
        "from numpy.random import choice, randint\n",
        "\n",
        "np.random.seed(0)\n",
        "pad, unk, sos, eos = 0, 1, 2, 3\n",
        "emb_size = 5  # 0, 1, 2, 3, 4\n",
        "batch_size = 3\n",
        "x1_lens = randint(1, 4, batch_size)\n",
        "print(x1_lens)\n",
        "max_len = np.max(x1_lens)\n",
        "print(max_len)\n",
        "x1_base = randint(1, 5, (batch_size, 4))\n",
        "bx1 = np.array([[2] + list(choice([1, 4], x)) + [3] + [0] * (max_len - x) for x in x1_lens])\n",
        "print(bx1)\n",
        "\n",
        "def predict(bx1, bx2, emb_size=5):\n",
        "  np.random.seed(0)\n",
        "  batch_size, len_x2 = bx2.shape\n",
        "  a = np.random.pareto(2, (len_x2, batch_size, emb_size))  # activations\n",
        "  a = a.swapaxes(0, 1)  # for consitency\n",
        "  s = np.exp(a)/np.sum(np.exp(a), axis=2, keepdims=True)  # softmax\n",
        "  by2p = np.log(s)  # log softmax\n",
        "  return by2p\n",
        "\n",
        "import torch\n",
        "from torch.distributions.pareto import Pareto\n",
        "import torch.nn.functional as F\n",
        "\n",
        "def torch_predict(bx1, bx2, emb_size=5):\n",
        "  torch.manual_seed(0)\n",
        "  batch_size, len_x2 = bx2.size()\n",
        "  bx1 = torch.repeat_interleave(bx1, int(batch_size/bx1.size(0)), 0)\n",
        "  scale, alpha = bx1[:,0], bx1[:,-1]\n",
        "  scale[scale == 0] = 1\n",
        "  alpha[alpha == 0] = 1\n",
        "  p = Pareto(scale.float(), alpha.float())  # Pareto\n",
        "  a = p.sample((len_x2, emb_size))\n",
        "  a = a.permute(2, 0, 1)\n",
        "  by2p = F.log_softmax(a, 2)  # log softmax predictions\n",
        "  return by2p\n",
        "\n",
        "bx1 = torch.tensor(bx1)\n",
        "bx2 = torch.full((batch_size, 1), 2)\n",
        "torch_predict(bx1, bx2)"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[1 2 1]\n",
            "2\n",
            "[[2 1 3 0]\n",
            " [2 1 1 3]\n",
            " [2 4 3 0]]\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "/pytorch/aten/src/ATen/native/TensorFactories.cpp:361: UserWarning: Deprecation warning: In a future PyTorch release torch.full will no longer return tensors of floating dtype by default. Instead, a bool fill_value will return a tensor of torch.bool dtype, and an integral fill_value will return a tensor of torch.long dtype. Set the optional `dtype` or `out` arguments to suppress this warning.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[[ 0.0000e+00, -4.1549e+01, -6.4350e+01, -6.3975e+01, -6.3683e+01]],\n",
              "\n",
              "        [[-1.2998e+00, -1.4892e+00, -2.0037e+00, -1.8685e+00, -1.5480e+00]],\n",
              "\n",
              "        [[-5.8917e+00, -1.5008e-02, -4.8140e+00, -7.1616e+00, -5.7313e+00]]])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 1
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KJItFZCZW5vF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def greedy(bx1, sos=2, eos=3, max_len=10):\n",
        "  bx2 = np.full((batch_size, 1), sos)  # batch with <sos> token\n",
        "  # stop when predictions len > max_len or all have <eos> token\n",
        "  while 1 + bx2.shape[1] < max_len and not np.all(np.any(bx2 == eos, axis=1)):\n",
        "    # print(bx2)\n",
        "    by2p = predict(bx1, bx2)\n",
        "    lp = by2p[:,-1,:]  # last prediction\n",
        "    # print(by2)\n",
        "    # print(lp)\n",
        "    next_bx2 = np.argmax(lp, axis=1).reshape((lp.shape[0], 1))\n",
        "    # print(next_bx2)\n",
        "    bx2 = np.concatenate((bx2, next_bx2), axis=1)\n",
        "    # print(bx2)\n",
        "    # print(bx2 == eos)\n",
        "    # print(np.all(np.any(bx2 == eos, axis=1)))\n",
        "  return by2p\n",
        "\n",
        "def greedy_many_to_one(bx1, sos=2, eos=3, max_len=10, reduction=np.sum):\n",
        "  bx2 = np.full((bx1.shape[0], 1), sos)  # batch with <sos> token\n",
        "  # stop when predictions len > max_len or all have <eos> token\n",
        "  while 1 + bx2.shape[1] < max_len and not np.all(np.any(bx2 == eos, axis=1)):\n",
        "    print(bx2)\n",
        "    by2p = predict(bx1, bx2)\n",
        "    lp = by2p[:,-1,:]  # last prediction\n",
        "    print(lp)\n",
        "    lp_red = reduction(lp, axis=0)\n",
        "    print(lp_red)\n",
        "    # print(by2)\n",
        "    # print(lp)\n",
        "    # next_bx2 = np.argmax(lp, axis=1).reshape((lp.shape[0], 1))\n",
        "    next_x2 = np.argmax(lp_red, axis=0)\n",
        "    # print(next_x2)\n",
        "    next_bx2 = np.full((lp.shape[0], 1), next_x2)\n",
        "    # print(next_bx2)\n",
        "    bx2 = np.concatenate((bx2, next_bx2), axis=1)\n",
        "    # print(bx2)\n",
        "    # print(bx2 == eos)\n",
        "    # print(np.all(np.any(bx2 == eos, axis=1)))\n",
        "  return by2p\n",
        "\n",
        "def forced(bx1, bx2):\n",
        "  by2p = predict(bx1, bx2)\n",
        "  return by2p\n",
        "\n",
        "def many_to_one(by2p, reduction=np.sum):\n",
        "  y2p = np.sum(by2p[:,:-1,:], axis=0)\n",
        "  y2 = np.argmax(y2p, axis=1)\n",
        "  return y2\n",
        "\n",
        "# by2p = greedy_many_to_one(bx1, reduction=np.sum)\n",
        "# print(by2p)\n",
        "# by2p = greedy_many_to_one(bx1, reduction=np.max)\n",
        "# print(by2p)\n",
        "# by2p = greedy_many_to_one(bx1, reduction=np.mean)\n",
        "# print(by2p)\n",
        "# y2 = np.argmax(np.sum(by2p[:,:-1,:], axis=0), axis=1)\n",
        "# print(y2)\n",
        "# print(many_to_one_sum(by2p))\n",
        "# by2p = greedy(bx1)\n",
        "# print(by2p)\n",
        "# by2 = np.argmax(by2p[:,:-1,:], axis=2)\n",
        "# print(by2)\n",
        "# print(many_to_one(by2p, np.sum))\n",
        "# print(many_to_one(by2p, np.max))\n",
        "# print(many_to_one(by2p, np.mean))\n",
        "# forced\n",
        "# bx2 = np.concatenate((np.full((batch_size, 1), 2), by2), axis=1)\n",
        "# print(bx2)\n",
        "# by2p = forced(bx1, bx2)\n",
        "# # print(bpy2)\n",
        "# by2 = np.argmax(by2p[:,:-1,:], axis=2)\n",
        "# print(by2)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rpAruRIeQYdY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# + do many_to_one search with batch_reduction\n",
        "# + try Viterbi algorithm for best beam searching (like beam search but in reversed direction)\n",
        "# TODO maybe try spectral beams? (With multiple reductions + Viterbi)\n",
        "# bx1 - encoder batch input\n",
        "# bx2 - decoder batch input\n",
        "# by2p - encoder batch output probabilties \n",
        "# by2 - encoder batch output"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "35vwuEvbsADZ",
        "colab_type": "code",
        "outputId": "08823761-0c6e-4b73-b250-c519bc05318e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 123
        }
      },
      "source": [
        "def beam(bx1, sos=2, eos=3, max_len=6,\n",
        "         beam_width=2, beam_depth=2, depth_reduction=np.sum):\n",
        "  batch_size = bx1.shape[0]\n",
        "  bx2 = np.full((batch_size, 1), sos)  # base batch with <sos> token\n",
        "  cur_len = bx2.shape[1] - 1  # without <sos>\n",
        "  # stop when cur_len > max_len or all predictions have <eos> token\n",
        "  while cur_len < max_len and not np.all(np.any(bx2 == eos, axis=1)):\n",
        "    bx2t = bx2.copy()  # temporal batch bx2\n",
        "    for i in range(beam_depth):\n",
        "      by2p = predict(bx1, bx2t)  # all predictions\n",
        "      lp = by2p[:,-1,:]  # last prediction\n",
        "      next_bx2 = np.argsort(-lp, axis=1)[:,:beam_width]  # top \"beam_width\" last predictions\n",
        "      bx2t = np.repeat(bx2t, beam_width, axis=0)  # multiply batch_size by beam_width\n",
        "      next_bx2 = next_bx2.reshape(batch_size*beam_width**(i+1), 1)  # to new batch_size\n",
        "      bx2t = np.concatenate((bx2t, next_bx2), axis=1)  # update batch\n",
        "    # Prediction with temporal batch\n",
        "    by2p = predict(bx1, bx2t)  # all predictions by last temporal batch\n",
        "    by2p = by2p[:,cur_len:-1]  # [base part:-last]  # remove old predictions\n",
        "    bx2t = bx2t[:,cur_len+1:]  # [base part + <sos>:]  # remove old labels\n",
        "    # Best beams searching\n",
        "    by2p = np.take_along_axis(by2p, bx2t[:,np.newaxis], axis=2)  # beams transitions predictions\n",
        "    by2p = by2p.reshape(batch_size, beam_width**beam_depth, -1)  # reshape to base batch shape (like)\n",
        "    by2p_red = depth_reduction(by2p, axis=2)  # beams reduction, i.e. giving them scores\n",
        "    best_beams = np.argmax(by2p_red, axis=1)  # get indices of the best beams\n",
        "    # Base batch updating\n",
        "    bx2t = bx2t.reshape(batch_size, beam_width**beam_depth, -1)  # reshape to base batch shape (like)\n",
        "    bx2t = np.take_along_axis(bx2t, best_beams[:, np.newaxis, np.newaxis], axis=1)  # get best beams\n",
        "    bx2t = bx2t.reshape(batch_size, -1)  # to new labels to base batch_size\n",
        "    bx2 = np.concatenate((bx2, bx2t), axis=1)  # update base batch\n",
        "    cur_len = bx2.shape[1] - 1  # without <sos>\n",
        "  if bx2.shape[1] - 1 > max_len:  # cut to max_len\n",
        "    bx2 = bx2[:max_len + 1]  # <sos> included\n",
        "  return bx2\n",
        "\n",
        "by2 = beam(bx1, depth_reduction=np.sum)\n",
        "print(by2)\n",
        "by2 = beam(bx1, depth_reduction=np.mean)\n",
        "print(by2)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[[2 2 4 3 3]\n",
            " [2 2 0 3 3]\n",
            " [2 3 2 4 3]]\n",
            "[[2 2 4 3 3]\n",
            " [2 2 0 3 3]\n",
            " [2 3 2 4 3]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mVi25jHphkjG",
        "colab_type": "code",
        "outputId": "e6d16538-8421-49a3-d53a-23d340bb097e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 176
        }
      },
      "source": [
        "def torch_beam(bx1, sos=2, eos=3, max_len=10,\n",
        "               beam_width=2, beam_depth=2, depth_reduction=torch.sum):\n",
        "  batch_size = bx1.size(0)  # input batch_size\n",
        "  bx2 = torch.full((batch_size, 1), sos, dtype=bx1.dtype)  # batch with <sos>\n",
        "  # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "  while bx2.size(1) - 1 < max_len and not torch.all(torch.any(bx2 == eos, axis=1)):\n",
        "    beam_scores = torch.empty((batch_size, 0))  # scores for each beam\n",
        "    bx1t = bx1.clone()\n",
        "    for i in range(beam_depth):                                           \n",
        "      new_batch_size = batch_size*beam_width**(i+1)\n",
        "      by2p = torch_predict(bx1t, bx2)  # predict\n",
        "      next_by2p, next_bx2 = torch.topk(by2p[:,-1,:], beam_width)  # beams to top k last predictions\n",
        "      next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "      next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "      beam_scores = torch.repeat_interleave(beam_scores, beam_width, 0)  # increase batch for new scores\n",
        "      bx2 = torch.repeat_interleave(bx2, beam_width, 0)  # increase batch for new beams\n",
        "      bx2 = torch.cat((bx2, next_bx2), 1)  # add beams\n",
        "      beam_scores = torch.cat((beam_scores, next_by2p), 1)  # add beams scores\n",
        "      bx1t = torch.repeat_interleave(bx1t, beam_width, 0)  # increase batch for new beams\n",
        "    beam_scores = depth_reduction(beam_scores, axis=1) # cumulative beams scores\n",
        "    beam_scores = beam_scores.view(batch_size, -1)  # split scores into batches\n",
        "    best_beams = torch.argmax(beam_scores, axis=1, keepdim=True)  # best beams\n",
        "    bx2 = bx2.view(batch_size, beam_width**beam_depth, -1)  # split beams into batches\n",
        "    # XXX its fucking magic... (return to input batch_size)\n",
        "    best_beams = best_beams.unsqueeze(2).expand(best_beams.size(0), best_beams.size(1), bx2.size(2))\n",
        "    bx2 = torch.gather(bx2, 1, best_beams)\n",
        "    bx2 = bx2.view(batch_size, -1)\n",
        "  return bx2\n",
        "\n",
        "# bx1 = torch.full((5, 1), 0, dtype=torch.long)\n",
        "bx1 = torch.randint(0, 4, (3, 3))\n",
        "print(bx1)\n",
        "by2 = torch_beam(bx1, max_len=10,\n",
        "                 beam_width=2, beam_depth=2, depth_reduction=torch.sum)\n",
        "print(by2)\n",
        "by2 = torch_beam(bx1, depth_reduction=torch.mean)\n",
        "print(by2)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[1, 0, 3],\n",
            "        [0, 3, 1],\n",
            "        [2, 3, 3]])\n",
            "tensor([[2, 0, 3, 0, 3, 1, 2, 4, 4, 1, 1],\n",
            "        [2, 0, 1, 0, 0, 2, 1, 0, 0, 4, 3],\n",
            "        [2, 1, 2, 1, 4, 3, 2, 3, 2, 1, 2]])\n",
            "tensor([[2, 0, 3, 0, 3, 1, 2, 4, 4, 1, 1],\n",
            "        [2, 0, 1, 0, 0, 2, 1, 0, 0, 4, 3],\n",
            "        [2, 1, 2, 1, 4, 3, 2, 3, 2, 1, 2]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "02e688d9-3d68-4133-baef-a38261d6d068",
        "id": "5otwtJ0V_0Do",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 283
        }
      },
      "source": [
        "def torch_beam_many_to_one(bx1, sos=2, eos=3, max_len=10,\n",
        "                           beam_width=2, beam_depth=2, \n",
        "                           depth_reduction=torch.sum,\n",
        "                           batch_reduction=torch.sum):\n",
        "  batch_size = bx1.size(0)  # input batch_size\n",
        "  bx2 = torch.full((batch_size, 1), sos, dtype=bx1.dtype)  # batch with <sos>\n",
        "  # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "  while bx2.size(1) - 1 < max_len and not torch.all(torch.any(bx2 == eos, axis=1)):\n",
        "    beam_scores = torch.empty((batch_size, 0))  # scores for each beam\n",
        "    bx1t = bx1.clone()\n",
        "    for i in range(beam_depth):\n",
        "      prev_batch_size = batch_size*beam_width**(i)\n",
        "      new_batch_size = batch_size*beam_width**(i+1)\n",
        "      by2p = torch_predict(bx1t, bx2)  # predict\n",
        "      # batch reduction\n",
        "      by2p = batch_reduction(by2p, axis=0) # cumulative batch predictions\n",
        "      next_by2p, next_bx2 = torch.topk(by2p[-1,:], beam_width)  # beams to top k last predictions\n",
        "      next_bx2 = next_bx2.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "      next_by2p = next_by2p.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "      # batch reduction   \n",
        "      next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "      next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "      beam_scores = torch.repeat_interleave(beam_scores, beam_width, 0)  # increase batch for new scores\n",
        "      bx2 = torch.repeat_interleave(bx2, beam_width, 0)  # increase batch for new beams\n",
        "      bx2 = torch.cat((bx2, next_bx2), 1)  # add beams\n",
        "      beam_scores = torch.cat((beam_scores, next_by2p), 1)  # add beams scores\n",
        "      bx1t = torch.repeat_interleave(bx1t, beam_width, 0)  # increase batch for new beams\n",
        "    beam_scores = depth_reduction(beam_scores, axis=1) # cumulative beams scores\n",
        "    best_beams = torch.argmax(beam_scores, axis=0, keepdim=True)  # best beams\n",
        "    bx2 = bx2[best_beams]  # best beam of all batches\n",
        "    bx2 = bx2.repeat(batch_size, 1)  # return to input batch size\n",
        "  return bx2\n",
        "\n",
        "# bx1 = torch.full((3, 1), 0, dtype=torch.long)\n",
        "bx1 = torch.randint(1, 4, (3, 3))\n",
        "print(bx1)\n",
        "bx2 = torch_beam_many_to_one(bx1, depth_reduction=torch.sum,\n",
        "                             batch_reduction=torch.sum)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_many_to_one(bx1, depth_reduction=torch.mean,\n",
        "                             batch_reduction=torch.sum)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_many_to_one(bx1, depth_reduction=torch.mean, \n",
        "                             batch_reduction=torch.mean)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_many_to_one(bx1, depth_reduction=torch.mean,\n",
        "                             batch_reduction=torch.mean)\n",
        "print(bx2)"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[2, 1, 3],\n",
            "        [3, 3, 3],\n",
            "        [1, 2, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WKGP6eMnQSPi",
        "colab_type": "code",
        "outputId": "20ff2a73-19a2-42ff-b6e2-4db1b27554ff",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 496
        }
      },
      "source": [
        "def torch_beam_reduction_many_to_one(bx1, sos=2, eos=3, max_len=10,\n",
        "                                     beam_width=2, beam_depth=2, \n",
        "                                     depth_reduction=torch.sum,\n",
        "                                     beam_reduction=torch.sum,\n",
        "                                     batch_reduction=torch.sum):\n",
        "  batch_size = bx1.size(0)  # input batch_size\n",
        "  bx2 = torch.full((batch_size, 1), sos, dtype=bx1.dtype)  # batch with <sos>\n",
        "  # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "  while bx2.size(1) - 1 < max_len and not torch.all(torch.any(bx2 == eos, axis=1)):\n",
        "    beam_scores = torch.empty((batch_size, 0))  # scores for each beam\n",
        "    bx1t = bx1.clone()\n",
        "    for i in range(beam_depth):\n",
        "      prev_batch_size = batch_size*beam_width**(i)\n",
        "      new_batch_size = batch_size*beam_width**(i+1)\n",
        "      by2p = torch_predict(bx1t, bx2)  # predict\n",
        "      # beam reduction\n",
        "      if i > 0:\n",
        "        prev_prev_batch_size = batch_size*beam_width**(i-1)\n",
        "        by2p = by2p.view(prev_prev_batch_size, beam_width, by2p.size(1), -1)  # split predictions into batches\n",
        "        by2p = beam_reduction(by2p, axis=1) # cumulative beams predictions\n",
        "      # batch reduction\n",
        "      by2p = batch_reduction(by2p, axis=0) # cumulative batch predictions\n",
        "      next_by2p, next_bx2 = torch.topk(by2p[-1,:], beam_width)  # beams to top k last predictions\n",
        "      next_bx2 = next_bx2.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "      next_by2p = next_by2p.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "      # basic beam\n",
        "      next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "      next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "      beam_scores = torch.repeat_interleave(beam_scores, beam_width, 0)  # increase batch for new scores\n",
        "      bx2 = torch.repeat_interleave(bx2, beam_width, 0)  # increase batch for new beams\n",
        "      bx2 = torch.cat((bx2, next_bx2), 1)  # add beams\n",
        "      beam_scores = torch.cat((beam_scores, next_by2p), 1)  # add beams scores\n",
        "      bx1t = torch.repeat_interleave(bx1t, beam_width, 0)  # increase batch for new beams\n",
        "    # best batch\n",
        "    beam_scores = depth_reduction(beam_scores, axis=1) # cumulative beams scores\n",
        "    best_beams = torch.argmax(beam_scores, axis=0, keepdim=True)  # best beams\n",
        "    bx2 = bx2[best_beams]  # best beam of all batches\n",
        "    bx2 = bx2.repeat(batch_size, 1)  # return to input batch size\n",
        "  return bx2\n",
        "\n",
        "# bx1 = torch.full((3, 1), 0, dtype=torch.long)\n",
        "bx1 = torch.randint(1, 4, (3, 3))\n",
        "print(bx1)\n",
        "bx2 = torch_beam_reduction_many_to_one(bx1,\n",
        "                                       depth_reduction=torch.sum,\n",
        "                                       beam_reduction=torch.sum,\n",
        "                                       batch_reduction=torch.sum)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_reduction_many_to_one(bx1, \n",
        "                                       depth_reduction=torch.mean,\n",
        "                                       beam_reduction=torch.sum,\n",
        "                                       batch_reduction=torch.sum)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_reduction_many_to_one(bx1, \n",
        "                                       depth_reduction=torch.mean, \n",
        "                                       beam_reduction=torch.sum,\n",
        "                                       batch_reduction=torch.mean)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_reduction_many_to_one(bx1, \n",
        "                                       depth_reduction=torch.mean,\n",
        "                                       beam_reduction=torch.sum,\n",
        "                                       batch_reduction=torch.mean)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_reduction_many_to_one(bx1,\n",
        "                                       depth_reduction=torch.sum,\n",
        "                                       beam_reduction=torch.mean,\n",
        "                                       batch_reduction=torch.sum)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_reduction_many_to_one(bx1, \n",
        "                                       depth_reduction=torch.mean,\n",
        "                                       beam_reduction=torch.mean,\n",
        "                                       batch_reduction=torch.sum)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_reduction_many_to_one(bx1, \n",
        "                                       depth_reduction=torch.mean, \n",
        "                                       beam_reduction=torch.mean,\n",
        "                                       batch_reduction=torch.mean)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_reduction_many_to_one(bx1, \n",
        "                                       depth_reduction=torch.mean,\n",
        "                                       beam_reduction=torch.mean,\n",
        "                                       batch_reduction=torch.mean)\n",
        "print(bx2)"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[1, 2, 2],\n",
            "        [2, 2, 3],\n",
            "        [2, 2, 2]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n",
            "tensor([[2, 0, 3],\n",
            "        [2, 0, 3],\n",
            "        [2, 0, 3]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OKqvYSzPVifr",
        "colab_type": "code",
        "outputId": "b28bab03-2350-47b6-e6ac-2589e78c0d3f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 283
        }
      },
      "source": [
        "def torch_beam_reduction(bx1, sos=2, eos=3, max_len=10, \n",
        "                         beam_width=2, beam_depth=2, \n",
        "                         beam_reduction=torch.sum,\n",
        "                         depth_reduction=torch.sum):\n",
        "  batch_size = bx1.size(0)  # input batch_size\n",
        "  bx2 = torch.full((batch_size, 1), sos, dtype=bx1.dtype)  # batch with <sos>\n",
        "  # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "  while bx2.size(1) - 1 < max_len and not torch.all(torch.any(bx2 == eos, axis=1)):\n",
        "    beam_scores = torch.empty((batch_size, 0))  # scores for each beam\n",
        "    bx1t = bx1.clone()\n",
        "    for i in range(beam_depth):\n",
        "      new_batch_size = batch_size*beam_width**(i+1)\n",
        "      by2p = torch_predict(bx1t, bx2)  # predict\n",
        "      if i > 0:  # beam reduction\n",
        "        prev_prev_batch_size = batch_size*beam_width**(i-1)\n",
        "        by2p = by2p.view(prev_prev_batch_size, beam_width, by2p.size(1), -1)  # split predictions into batches\n",
        "        by2p = beam_reduction(by2p, axis=1) # cumulative beams predictions\n",
        "        by2p = by2p.repeat(beam_width, 1, 1)  # return to prev batch size\n",
        "      next_by2p, next_bx2 = torch.topk(by2p[:,-1,:], beam_width)  # beams to top k last predictions\n",
        "      next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "      next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "      beam_scores = torch.repeat_interleave(beam_scores, beam_width, 0)  # increase batch for new scores\n",
        "      bx2 = torch.repeat_interleave(bx2, beam_width, 0)  # increase batch for new beams\n",
        "      bx2 = torch.cat((bx2, next_bx2), 1)  # add beams\n",
        "      beam_scores = torch.cat((beam_scores, next_by2p), 1)  # add beams scores\n",
        "      bx1t = torch.repeat_interleave(bx1t, beam_width, 0)  # increase batch for new beams\n",
        "    beam_scores = depth_reduction(beam_scores, axis=1) # cumulative beams scores\n",
        "    beam_scores = beam_scores.view(batch_size, -1)  # split scores into batches\n",
        "    best_beams = torch.argmax(beam_scores, axis=1, keepdim=True)  # best beams\n",
        "    bx2 = bx2.view(batch_size, beam_width**beam_depth, -1)  # split beams into batches\n",
        "    # XXX its fucking magic... (return to input batch_size)\n",
        "    best_beams = best_beams.unsqueeze(2).expand(best_beams.size(0), best_beams.size(1), bx2.size(2))\n",
        "    bx2 = torch.gather(bx2, 1, best_beams)\n",
        "    bx2 = bx2.view(batch_size, -1)\n",
        "  return bx2\n",
        "\n",
        "# bx1 = torch.full((5, 1), 0, dtype=torch.long)\n",
        "bx1 = torch.randint(0, 4, (3, 3))\n",
        "print(bx1)\n",
        "bx2 = torch_beam_reduction(bx1, depth_reduction=torch.sum, beam_reduction=torch.sum)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_reduction(bx1, depth_reduction=torch.mean, beam_reduction=torch.sum)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_reduction(bx1, depth_reduction=torch.sum, beam_reduction=torch.mean)\n",
        "print(bx2)\n",
        "bx2 = torch_beam_reduction(bx1, depth_reduction=torch.mean, beam_reduction=torch.mean)\n",
        "print(bx2)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[3, 3, 1],\n",
            "        [1, 1, 0],\n",
            "        [1, 1, 1]])\n",
            "tensor([[2, 0, 3, 4, 2, 1, 0],\n",
            "        [2, 1, 3, 0, 1, 3, 0],\n",
            "        [2, 1, 1, 1, 2, 3, 1]])\n",
            "tensor([[2, 0, 3, 4, 2, 1, 0],\n",
            "        [2, 1, 3, 0, 1, 3, 0],\n",
            "        [2, 1, 1, 1, 2, 3, 1]])\n",
            "tensor([[2, 0, 3, 0, 3, 1, 0],\n",
            "        [2, 1, 3, 0, 1, 2, 2],\n",
            "        [2, 1, 1, 1, 2, 3, 1]])\n",
            "tensor([[2, 0, 3, 0, 3, 1, 0],\n",
            "        [2, 1, 3, 0, 1, 2, 2],\n",
            "        [2, 1, 1, 1, 2, 3, 1]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "siMsdBerQBd7",
        "colab_type": "code",
        "outputId": "29f5913f-dc62-436c-94f5-4f1ff73dbbee",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 123
        }
      },
      "source": [
        "def torch_beam_viterbi(bx1, sos=2, eos=3, max_len=4,\n",
        "                       beam_width=2, beam_depth=2):\n",
        "  batch_size = bx1.size(0)  # input batch_size\n",
        "  bx2 = torch.full((batch_size, 1), sos, dtype=bx1.dtype)  # batch with <sos>\n",
        "  # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "  while bx2.size(1) - 1 < max_len and not torch.all(torch.any(bx2 == eos, axis=1)):\n",
        "    fwd_scores = torch.empty((batch_size, 0))  # scores for each beam\n",
        "    fwd_bx2 = bx2.clone()\n",
        "    bx1t = bx1.clone()\n",
        "    for i in range(beam_depth):\n",
        "      new_batch_size = batch_size*beam_width**(i+1)\n",
        "      by2p = torch_predict(bx1t, fwd_bx2)  # predict\n",
        "      next_by2p, next_bx2 = torch.topk(by2p[:,-1,:], beam_width)  # beams to top k last predictions\n",
        "      next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "      next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "      fwd_scores = torch.repeat_interleave(fwd_scores, beam_width, 0)  # increase batch for new scores\n",
        "      fwd_bx2 = torch.repeat_interleave(fwd_bx2, beam_width, 0)  # increase batch for new beams\n",
        "      fwd_bx2 = torch.cat((fwd_bx2, next_bx2), 1)  # add beams\n",
        "      fwd_scores = torch.cat((fwd_scores, next_by2p), 1)  # add beams scores\n",
        "      bx1t = torch.repeat_interleave(bx1t, beam_width, 0)  # increase batch for new beams\n",
        "    fwd_scores = fwd_scores.view(batch_size, beam_width**beam_depth, -1)  # split scores into batches\n",
        "    fwd_bx2 = fwd_bx2.view(batch_size, beam_width**beam_depth, -1)  # split beams into batches\n",
        "    bkw_bx2 = torch.empty((batch_size, 0), dtype=bx1.dtype)  # scores for each beam\n",
        "    bkw_scores =  torch.empty((batch_size, 0))  # scores for each beam\n",
        "    mask = torch.full((fwd_bx2.size(0), fwd_bx2.size(1), 1), False, dtype=torch.bool)\n",
        "    for i in range(fwd_scores.size(2)):\n",
        "      cur_values = fwd_bx2[:,:,-i-1:-i if i > 0 else None]\n",
        "      cur_scores = fwd_scores[:,:,-i-1:-i if i > 0 else None]\n",
        "      cur_scores[mask] = float('-inf')  # mask from prev step\n",
        "      best_scores, best_beams = torch.max(cur_scores, axis=1, keepdim=True)\n",
        "      best_values = torch.gather(cur_values, 1, best_beams).view(batch_size, -1)\n",
        "      exp_best_values = best_values[:,None,:].expand(-1, cur_values.size(1), -1)\n",
        "      mask = cur_values != exp_best_values\n",
        "      best_values = best_values.view(batch_size, -1)\n",
        "      bkw_bx2 = torch.cat((best_values, bkw_bx2), 1)\n",
        "      best_scores = best_scores.view(batch_size, -1)\n",
        "      bkw_scores = torch.cat((best_scores, bkw_scores), 1)\n",
        "    bx2 = torch.cat((bx2, bkw_bx2), 1)  # <sos> + back_bx2\n",
        "  return bx2\n",
        "\n",
        "bx1 = torch.full((5, 1), 0, dtype=torch.long)\n",
        "bx1 = torch.randint(0, 4, (3, 3))\n",
        "print(bx1)\n",
        "by2 = torch_beam_viterbi(bx1)\n",
        "print(by2)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[0, 3, 0],\n",
            "        [3, 0, 0],\n",
            "        [0, 0, 1]])\n",
            "tensor([[2, 0, 3, 4, 0],\n",
            "        [2, 1, 4, 0, 0],\n",
            "        [2, 2, 3, 2, 1]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4zu0Q4n6CEq2",
        "colab_type": "code",
        "outputId": "f8a9a995-f29a-43b6-c6b9-47f4396ac885",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 176
        }
      },
      "source": [
        "def torch_beam_bkw(bx1, sos=2, eos=3, max_len=4,\n",
        "                   beam_width=2, beam_depth=2, \n",
        "                   bkw_beam_width=2,\n",
        "                   depth_reduction=torch.sum):\n",
        "  batch_size = bx1.size(0)  # input batch_size\n",
        "  bx2 = torch.full((batch_size, 1), sos, dtype=bx1.dtype)  # batch with <sos>\n",
        "  # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "  while bx2.size(1) - 1 < max_len and not torch.all(torch.any(bx2 == eos, axis=1)):\n",
        "    fwd_scores = torch.empty((batch_size, 0))  # scores for each beam\n",
        "    fwd_bx2 = bx2.clone()\n",
        "    bx1t = bx1.clone()\n",
        "    for i in range(beam_depth):\n",
        "      new_batch_size = batch_size*beam_width**(i+1)\n",
        "      by2p = torch_predict(bx1t, fwd_bx2)  # predict\n",
        "      next_by2p, next_bx2 = torch.topk(by2p[:,-1,:], beam_width)  # beams to top k last predictions\n",
        "      next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "      next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "      fwd_scores = torch.repeat_interleave(fwd_scores, beam_width, 0)  # increase batch for new scores\n",
        "      fwd_bx2 = torch.repeat_interleave(fwd_bx2, beam_width, 0)  # increase batch for new beams\n",
        "      fwd_bx2 = torch.cat((fwd_bx2, next_bx2), 1)  # add beams\n",
        "      fwd_scores = torch.cat((fwd_scores, next_by2p), 1)  # add beams scores\n",
        "      bx1t = torch.repeat_interleave(bx1t, beam_width, 0)  # increase batch for new beams\n",
        "    fwd_scores = fwd_scores.view(batch_size, beam_width**beam_depth, -1)  # split scores into batches\n",
        "    fwd_bx2 = fwd_bx2.view(batch_size, beam_width**beam_depth, -1)  # split beams into batches\n",
        "    bkw_bx2 = torch.empty((batch_size, beam_width**beam_depth, 0), dtype=bx1.dtype)  # scores for each beam\n",
        "    bkw_scores = torch.empty((batch_size, beam_width**beam_depth, 0))  # scores for each beam\n",
        "    mask = torch.full((fwd_bx2.size(0), fwd_bx2.size(1), 1), False, dtype=torch.bool)\n",
        "    n_beams_per_batch = beam_width**beam_depth\n",
        "    for i in range(fwd_scores.size(2)):\n",
        "      old_n_beams_per_batch = n_beams_per_batch*bkw_beam_width**(i)\n",
        "      new_n_beams_per_batch = n_beams_per_batch*bkw_beam_width**(i+1)\n",
        "      cur_values = fwd_bx2[:,:,-i-1:-i if i > 0 else None]\n",
        "      cur_values = torch.repeat_interleave(cur_values, bkw_beam_width**i, 1)\n",
        "      cur_scores = fwd_scores[:,:,-i-1:-i if i > 0 else None]\n",
        "      cur_scores = torch.repeat_interleave(cur_scores, bkw_beam_width**i, 1)\n",
        "      cur_scores[mask] = float('-inf')  # mask from prev step\n",
        "      best_scores, best_beams = torch.topk(cur_scores, bkw_beam_width, 1)  # best beams\n",
        "      best_values = torch.gather(cur_values, 1, best_beams).view(batch_size, -1)\n",
        "      best_values = best_values.view(best_values.size(0), bkw_beam_width, -1)\n",
        "      best_values = torch.repeat_interleave(best_values, old_n_beams_per_batch, 1)\n",
        "      cur_values = cur_values.repeat(1, bkw_beam_width, 1)\n",
        "      mask = cur_values != best_values\n",
        "      bkw_bx2 = bkw_bx2.repeat(1, bkw_beam_width, 1)  # increase batch for new beams\n",
        "      bkw_bx2 = torch.cat((best_values, bkw_bx2), 2)\n",
        "      best_scores = torch.repeat_interleave(best_scores, old_n_beams_per_batch, 1)\n",
        "      bkw_scores = bkw_scores.repeat(1, bkw_beam_width, 1)  # increase batch for new beams\n",
        "      bkw_scores = torch.cat((best_scores, bkw_scores), 2)\n",
        "    beam_scores = depth_reduction(bkw_scores, axis=2) # cumulative beams scores\n",
        "    best_beams = torch.argmax(beam_scores, axis=1, keepdim=True)  # best beams\n",
        "    best_beams = best_beams.unsqueeze(2).expand(best_beams.size(0), best_beams.size(1), bkw_bx2.size(2))\n",
        "    bkw_bx2 = torch.gather(bkw_bx2, 1, best_beams) # best beam of all batches\n",
        "    bkw_bx2 = bkw_bx2.view(batch_size, -1)\n",
        "    bx2 = torch.cat((bx2, bkw_bx2), 1)  # <sos> + bkw_bx2\n",
        "  return bx2\n",
        "\n",
        "bx1 = torch.full((5, 1), 0, dtype=torch.long)\n",
        "bx1 = torch.randint(0, 4, (3, 3))\n",
        "print(bx1)\n",
        "by2 = torch_beam_bkw(bx1, depth_reduction=torch.sum)\n",
        "print(by2)\n",
        "by2 = torch_beam_bkw(bx1, depth_reduction=torch.mean)\n",
        "print(by2)"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[0, 1, 2],\n",
            "        [1, 1, 3],\n",
            "        [1, 0, 0]])\n",
            "tensor([[2, 0, 3, 0, 0],\n",
            "        [2, 0, 4, 0, 2],\n",
            "        [2, 1, 3, 1, 1]])\n",
            "tensor([[2, 0, 3, 0, 0],\n",
            "        [2, 0, 4, 0, 2],\n",
            "        [2, 1, 3, 1, 1]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ygjkg48HzbFs",
        "colab_type": "code",
        "outputId": "81362a67-0e12-4e5a-bb1a-5ba50fc4fa3b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 283
        }
      },
      "source": [
        "def torch_beam_bkw_many_to_one(bx1, sos=2, eos=3, max_len=10,\n",
        "                               beam_width=3, beam_depth=2, \n",
        "                               bkw_beam_width=2,\n",
        "                               depth_reduction=torch.sum,\n",
        "                               batch_reduction=torch.sum):\n",
        "  batch_size = bx1.size(0)  # input batch_size\n",
        "  bx2 = torch.full((batch_size, 1), sos, dtype=bx1.dtype)  # batch with <sos>\n",
        "  # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "  while bx2.size(1) - 1 < max_len and not torch.all(torch.any(bx2 == eos, axis=1)):\n",
        "    fwd_scores = torch.empty((batch_size, 0))  # scores for each beam\n",
        "    fwd_bx2 = bx2.clone()\n",
        "    bx1t = bx1.clone()\n",
        "    for i in range(beam_depth):\n",
        "      prev_batch_size = batch_size*beam_width**(i)\n",
        "      new_batch_size = batch_size*beam_width**(i+1)\n",
        "      by2p = torch_predict(bx1t, fwd_bx2)  # predict\n",
        "      # fwd batch reduction\n",
        "      by2p = batch_reduction(by2p, axis=0) # cumulative batch predictions\n",
        "      next_by2p, next_bx2 = torch.topk(by2p[-1,:], beam_width)  # beams to top k last predictions\n",
        "      next_bx2 = next_bx2.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "      next_by2p = next_by2p.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "      # fwd\n",
        "      next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "      next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "      fwd_scores = torch.repeat_interleave(fwd_scores, beam_width, 0)  # increase batch for new scores\n",
        "      fwd_bx2 = torch.repeat_interleave(fwd_bx2, beam_width, 0)  # increase batch for new beams\n",
        "      fwd_bx2 = torch.cat((fwd_bx2, next_bx2), 1)  # add beams\n",
        "      fwd_scores = torch.cat((fwd_scores, next_by2p), 1)  # add beams scores\n",
        "      bx1t = torch.repeat_interleave(bx1t, beam_width, 0)  # increase batch for new beams\n",
        "    # bkw\n",
        "    bkw_bx2 = torch.empty((new_batch_size, 0), dtype=bx1.dtype)  # scores for each beam\n",
        "    bkw_scores = torch.empty((new_batch_size, 0))  # scores for each beam\n",
        "    mask = torch.full((new_batch_size, 1), False, dtype=torch.bool)\n",
        "    fwd_batch_size = new_batch_size\n",
        "    for i in range(fwd_scores.size(1)):\n",
        "      prev_batch_size = fwd_batch_size*bkw_beam_width**(i)\n",
        "      new_batch_size = fwd_batch_size*bkw_beam_width**(i+1)\n",
        "      cur_values = fwd_bx2[:,-i-1:-i if i > 0 else None]\n",
        "      cur_values = torch.repeat_interleave(cur_values, bkw_beam_width**i, 0)\n",
        "      cur_scores = fwd_scores[:,-i-1:-i if i > 0 else None]\n",
        "      cur_scores = torch.repeat_interleave(cur_scores, bkw_beam_width**i, 0)\n",
        "      cur_scores[mask] = float('-inf')  # mask from prev step\n",
        "      best_scores, best_beams = torch.topk(cur_scores, bkw_beam_width, 0)  # best beams\n",
        "      best_values = cur_values[best_beams]\n",
        "      mask = cur_values != best_values\n",
        "      mask = mask.view(new_batch_size, -1)\n",
        "      best_values = best_values.view(bkw_beam_width, -1)\n",
        "      best_values = torch.repeat_interleave(best_values, prev_batch_size, 0)\n",
        "      bkw_bx2 = bkw_bx2.repeat(bkw_beam_width, 1)  # increase batch for new beams\n",
        "      bkw_bx2 = torch.cat((best_values, bkw_bx2), 1)\n",
        "      best_scores = torch.repeat_interleave(best_scores, prev_batch_size, 0)\n",
        "      bkw_scores = bkw_scores.repeat(bkw_beam_width, 1)  # increase batch for new beams\n",
        "      bkw_scores = torch.cat((best_scores, bkw_scores), 1)\n",
        "    # bkw batch reduction\n",
        "    bkw_scores = depth_reduction(bkw_scores, axis=1) # cumulative beams scores\n",
        "    best_beams = torch.argmax(bkw_scores, axis=0, keepdim=True)  # best beams\n",
        "    bkw_bx2 = bkw_bx2[best_beams]  # best beam of all batches\n",
        "    bkw_bx2 = bkw_bx2.repeat(batch_size, 1)  # return to input batch size\n",
        "    bx2 = torch.cat((bx2, bkw_bx2), 1)  # <sos> + bkw_bx2\n",
        "  return bx2\n",
        "\n",
        "bx1 = torch.full((20, 1), 0, dtype=torch.long)\n",
        "bx1 = torch.randint(0, 4, (3, 3))\n",
        "print(bx1)\n",
        "by2 = torch_beam_bkw_many_to_one(bx1, depth_reduction=torch.sum, \n",
        "                                 batch_reduction=torch.sum)\n",
        "print(by2)\n",
        "by2 = torch_beam_bkw_many_to_one(bx1, depth_reduction=torch.mean, \n",
        "                                 batch_reduction=torch.sum)\n",
        "print(by2)\n",
        "by2 = torch_beam_bkw_many_to_one(bx1, depth_reduction=torch.sum, \n",
        "                                 batch_reduction=torch.mean)\n",
        "print(by2)\n",
        "by2 = torch_beam_bkw_many_to_one(bx1, depth_reduction=torch.mean, \n",
        "                                 batch_reduction=torch.mean)\n",
        "print(by2)"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[0, 3, 0],\n",
            "        [3, 0, 0],\n",
            "        [0, 0, 1]])\n",
            "tensor([[2, 0, 0, 1, 3],\n",
            "        [2, 0, 0, 1, 3],\n",
            "        [2, 0, 0, 1, 3]])\n",
            "tensor([[2, 0, 0, 1, 3],\n",
            "        [2, 0, 0, 1, 3],\n",
            "        [2, 0, 0, 1, 3]])\n",
            "tensor([[2, 0, 0, 1, 3],\n",
            "        [2, 0, 0, 1, 3],\n",
            "        [2, 0, 0, 1, 3]])\n",
            "tensor([[2, 0, 0, 1, 3],\n",
            "        [2, 0, 0, 1, 3],\n",
            "        [2, 0, 0, 1, 3]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cfgwGkcxqzZe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# import torch\n",
        "# bt = bx1[:,1:]  # without <sos>\n",
        "# print(bt)\n",
        "# print(bt[bt != 0])\n",
        "# bp = by2[:,:-1,:]  # without last prediction\n",
        "# print(bp)\n",
        "# print(bt != 0)\n",
        "# print(bp[bt != 0])\n",
        "\n",
        "# bp, bt = torch.tensor(bp), torch.tensor(bt)\n",
        "# loss = torch.nn.NLLLoss()\n",
        "# print(loss(bp[bt != 0], bt[bt != 0]))\n",
        "\n",
        "# loss = torch.nn.NLLLoss(ignore_index=0)\n",
        "# print(loss(torch.flatten(bp, 0, 1), torch.flatten(bt, 0)))\n",
        "\n",
        "# loss = 0\n",
        "# cnt = 0\n",
        "# for i, t in enumerate(bt):\n",
        "#   for j, l in enumerate(t):\n",
        "#     print(l)\n",
        "#     print(bp[i,j,l])\n",
        "#     loss += bp[i,j,l].item()\n",
        "#     cnt += 1\n",
        "# print(loss, loss/cnt)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1M95ua8Oa-BZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "reduction_map = {\n",
        "  'sum': torch.sum,\n",
        "  'mean': torch.mean,\n",
        "  'none': None\n",
        "}\n",
        "\n",
        "class Forced():\n",
        "  def __init__(self, **kwargs):\n",
        "    pass\n",
        "\n",
        "  def __call__(self, model, bx1, bx2):\n",
        "    by2p, attn = model(bx1, bx2)\n",
        "    y2 = torch.argmax(by2p[:,:-1,:], axis=2)\n",
        "    return y2, attn  \n",
        "\n",
        "class Greedy():\n",
        "  def __init__(self, sos=2, eos=3, max_len=30, **kwargs):\n",
        "    self.sos, self.eos, self.max_len = sos, eos, max_len\n",
        "\n",
        "  def __call__(self, model, bx1, bx2=None):\n",
        "    batch_size = bx1.size(0)  # input batch_size\n",
        "    bx2 = torch.full((batch_size, 1), self.sos, dtype=bx1.dtype)  # batch with <sos>\n",
        "    # stop when predictions len > max_len or all have <eos> token\n",
        "    while bx2.size(1) - 1 < self.max_len and not torch.all(torch.any(bx2 == self.eos, axis=1)):\n",
        "      by2p, attn = model(bx1, bx2)\n",
        "      next_bx2 = torch.argmax(by2p[:,-1,:], axis=1, keepdim=True)\n",
        "      bx2 = torch.cat((bx2, next_bx2), 1)\n",
        "    return bx2[:,1:], attn  \n",
        "\n",
        "class Beam():\n",
        "  def __init__(self, sos=2, eos=3, max_len=30, beam_width=2, beam_depth=2,\n",
        "               depth_reduction='sum', beam_reduction='none', \n",
        "               batch_reduction='none', **kwargs):\n",
        "    self.sos, self.eos, self.max_len = sos, eos, max_len\n",
        "    self.beam_width, self.beam_depth = beam_width, beam_depth\n",
        "    self.depth_reduction = reduction_map[depth_reduction]\n",
        "    self.beam_reduction = reduction_map[beam_reduction]\n",
        "    self.batch_reduction = reduction_map[batch_reduction]\n",
        "\n",
        "  def __call__(self, model, bx1, bx2=None):\n",
        "    batch_size = bx1.size(0)  # input batch_size\n",
        "    bx2 = torch.full((batch_size, 1), self.sos, dtype=bx1.dtype)  # batch with <sos>\n",
        "    # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "    while bx2.size(1) - 1 < self.max_len and not torch.all(torch.any(bx2 == self.eos, axis=1)):\n",
        "      beam_scores = torch.empty((batch_size, 0))  # scores for each beam\n",
        "      bx1t = bx1.clone()\n",
        "      for i in range(self.beam_depth):                                           \n",
        "        new_batch_size = batch_size*self.beam_width**(i+1)\n",
        "        by2p, attn = model(bx1t, bx2)  # predict\n",
        "        if i > 0 and self.beam_reduction is not None:  # beam reduction\n",
        "          prev_prev_batch_size = batch_size*self.beam_width**(i-1)\n",
        "          by2p = by2p.view(prev_prev_batch_size, self.beam_width, by2p.size(1), -1)  # split predictions into batches\n",
        "          by2p = self.beam_reduction(by2p, axis=1) # cumulative beams predictions\n",
        "          by2p = by2p.repeat(self.beam_width, 1, 1)  # return to prev batch size\n",
        "        if self.batch_reduction is not None:  # batch reduction\n",
        "          prev_batch_size = batch_size*self.beam_width**(i)\n",
        "          by2p = self.batch_reduction(by2p, axis=0) # cumulative batch predictions\n",
        "          next_by2p, next_bx2 = torch.topk(by2p[-1,:], self.beam_width)  # beams to top k last predictions\n",
        "          next_bx2 = next_bx2.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "          next_by2p = next_by2p.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "        else:\n",
        "          next_by2p, next_bx2 = torch.topk(by2p[:,-1,:], self.beam_width)  # beams to top k last predictions\n",
        "        next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "        next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "        beam_scores = torch.repeat_interleave(beam_scores, self.beam_width, 0)  # increase batch for new scores\n",
        "        bx2 = torch.repeat_interleave(bx2, self.beam_width, 0)  # increase batch for new beams\n",
        "        bx2 = torch.cat((bx2, next_bx2), 1)  # add beams\n",
        "        beam_scores = torch.cat((beam_scores, next_by2p), 1)  # add beams scores\n",
        "        bx1t = torch.repeat_interleave(bx1t, self.beam_width, 0)  # increase batch for new beams\n",
        "      beam_scores = self.depth_reduction(beam_scores, axis=1) # cumulative beams scores\n",
        "      if self.batch_reduction is None:\n",
        "        beam_scores = beam_scores.view(batch_size, -1)  # split scores into batches\n",
        "        best_beams = torch.argmax(beam_scores, axis=1, keepdim=True)  # best beams\n",
        "        bx2 = bx2.view(batch_size, self.beam_width**self.beam_depth, -1)  # split beams into batches\n",
        "        # XXX its fucking magic... (return to input batch_size)\n",
        "        best_beams = best_beams.unsqueeze(2).expand(best_beams.size(0), \n",
        "                                                    best_beams.size(1),\n",
        "                                                    bx2.size(2))\n",
        "        bx2 = torch.gather(bx2, 1, best_beams)\n",
        "        bx2 = bx2.view(batch_size, -1)\n",
        "      else:\n",
        "        best_beams = torch.argmax(beam_scores, axis=0, keepdim=True)  # best beams\n",
        "        bx2 = bx2[best_beams]  # best beam of all batches\n",
        "        bx2 = bx2.repeat(batch_size, 1)  # return to input batch size\n",
        "    return bx2[:,1:], attn\n",
        "\n",
        "class BeamBkw():\n",
        "  def __init__(self, sos=2, eos=3, max_len=30, beam_width=2, beam_depth=2,\n",
        "               depth_reduction='sum', beam_reduction='none', \n",
        "               batch_reduction='none', bkw_beam_width=2, **kwargs):\n",
        "    self.sos, self.eos, self.max_len = sos, eos, max_len\n",
        "    self.beam_width, self.beam_depth = beam_width, beam_depth\n",
        "    self.bkw_beam_width = bkw_beam_width\n",
        "    self.depth_reduction = reduction_map[depth_reduction]\n",
        "    self.beam_reduction = reduction_map[beam_reduction]\n",
        "    self.batch_reduction = reduction_map[batch_reduction]\n",
        "\n",
        "  def __call__(self, model, bx1, bx2=None):\n",
        "    batch_size = bx1.size(0)  # input batch_size\n",
        "    bx2 = torch.full((batch_size, 1), self.sos, dtype=bx1.dtype)  # batch with <sos>\n",
        "    # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "    while bx2.size(1) - 1 < self.max_len and not torch.all(torch.any(bx2 == self.eos, axis=1)):\n",
        "      fwd_scores = torch.empty((batch_size, 0))  # scores for each beam\n",
        "      fwd_bx2 = bx2.clone()\n",
        "      bx1t = bx1.clone()\n",
        "      for i in range(self.beam_depth):\n",
        "        new_batch_size = batch_size*self.beam_width**(i+1)\n",
        "        by2p, attn = model(bx1t, fwd_bx2)  # predict\n",
        "        if i > 0 and self.beam_reduction is not None:  # beam reduction\n",
        "          prev_prev_batch_size = batch_size*self.beam_width**(i-1)\n",
        "          by2p = by2p.view(prev_prev_batch_size, self.beam_width, by2p.size(1), -1)  # split predictions into batches\n",
        "          by2p = self.beam_reduction(by2p, axis=1) # cumulative beams predictions\n",
        "          by2p = by2p.repeat(self.beam_width, 1, 1)  # return to prev batch size\n",
        "        if self.batch_reduction is not None:\n",
        "          prev_batch_size = batch_size*self.beam_width**(i)\n",
        "          by2p = self.batch_reduction(by2p, axis=0) # cumulative batch predictions\n",
        "          next_by2p, next_bx2 = torch.topk(by2p[-1,:], self.beam_width)  # beams to top k last predictions\n",
        "          next_bx2 = next_bx2.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "          next_by2p = next_by2p.repeat(prev_batch_size, 1)  # return to prev batch size        \n",
        "        else:\n",
        "          next_by2p, next_bx2 = torch.topk(by2p[:,-1,:], self.beam_width)  # beams to top k last predictions        \n",
        "        next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "        next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "        fwd_scores = torch.repeat_interleave(fwd_scores, self.beam_width, 0)  # increase batch for new scores\n",
        "        fwd_bx2 = torch.repeat_interleave(fwd_bx2, self.beam_width, 0)  # increase batch for new beams\n",
        "        fwd_bx2 = torch.cat((fwd_bx2, next_bx2), 1)  # add beams\n",
        "        fwd_scores = torch.cat((fwd_scores, next_by2p), 1)  # add beams scores\n",
        "        bx1t = torch.repeat_interleave(bx1t, self.beam_width, 0)  # increase batch for new beams\n",
        "      if self.batch_reduction is None:\n",
        "        fwd_scores = fwd_scores.view(batch_size, self.beam_width**self.beam_depth, -1)  # split scores into batches\n",
        "        fwd_bx2 = fwd_bx2.view(batch_size, self.beam_width**self.beam_depth, -1)  # split beams into batches\n",
        "        bkw_bx2 = torch.empty((batch_size, self.beam_width**self.beam_depth, 0), dtype=bx1.dtype)  # scores for each beam\n",
        "        bkw_scores = torch.empty((batch_size, self.beam_width**self.beam_depth, 0))  # scores for each beam\n",
        "        mask = torch.full((fwd_bx2.size(0), fwd_bx2.size(1), 1), False, dtype=torch.bool)\n",
        "        n_beams_per_batch = self.beam_width**self.beam_depth\n",
        "        for i in range(fwd_scores.size(2)):\n",
        "          old_n_beams_per_batch = n_beams_per_batch*self.bkw_beam_width**(i)\n",
        "          new_n_beams_per_batch = n_beams_per_batch*self.bkw_beam_width**(i+1)\n",
        "          cur_values = fwd_bx2[:,:,-i-1:-i if i > 0 else None]\n",
        "          cur_values = torch.repeat_interleave(cur_values, self.bkw_beam_width**i, 1)\n",
        "          cur_scores = fwd_scores[:,:,-i-1:-i if i > 0 else None]\n",
        "          cur_scores = torch.repeat_interleave(cur_scores, self.bkw_beam_width**i, 1)\n",
        "          cur_scores[mask] = float('-inf')  # mask from prev step\n",
        "          best_scores, best_beams = torch.topk(cur_scores, self.bkw_beam_width, 1)  # best beams\n",
        "          best_values = torch.gather(cur_values, 1, best_beams).view(batch_size, -1)\n",
        "          best_values = best_values.view(best_values.size(0), self.bkw_beam_width, -1)\n",
        "          best_values = torch.repeat_interleave(best_values, old_n_beams_per_batch, 1)\n",
        "          cur_values = cur_values.repeat(1, self.bkw_beam_width, 1)\n",
        "          mask = cur_values != best_values\n",
        "          bkw_bx2 = bkw_bx2.repeat(1, self.bkw_beam_width, 1)  # increase batch for new beams\n",
        "          bkw_bx2 = torch.cat((best_values, bkw_bx2), 2)\n",
        "          best_scores = torch.repeat_interleave(best_scores, old_n_beams_per_batch, 1)\n",
        "          bkw_scores = bkw_scores.repeat(1, self.bkw_beam_width, 1)  # increase batch for new beams\n",
        "          bkw_scores = torch.cat((best_scores, bkw_scores), 2)\n",
        "        beam_scores = self.depth_reduction(bkw_scores, axis=2) # cumulative beams scores\n",
        "        best_beams = torch.argmax(beam_scores, axis=1, keepdim=True)  # best beams\n",
        "        best_beams = best_beams.unsqueeze(2).expand(best_beams.size(0), best_beams.size(1), bkw_bx2.size(2))\n",
        "        bkw_bx2 = torch.gather(bkw_bx2, 1, best_beams) # best beam of all batches\n",
        "        bkw_bx2 = bkw_bx2.view(batch_size, -1)\n",
        "        bx2 = torch.cat((bx2, bkw_bx2), 1)  # <sos> + bkw_bx2\n",
        "      else:  # batch reduction\n",
        "        bkw_bx2 = torch.empty((new_batch_size, 0), dtype=bx1.dtype)  # scores for each beam\n",
        "        bkw_scores = torch.empty((new_batch_size, 0))  # scores for each beam\n",
        "        mask = torch.full((new_batch_size, 1), False, dtype=torch.bool)\n",
        "        fwd_batch_size = new_batch_size\n",
        "        for i in range(fwd_scores.size(1)):\n",
        "          prev_batch_size = fwd_batch_size*self.bkw_beam_width**(i)\n",
        "          new_batch_size = fwd_batch_size*self.bkw_beam_width**(i+1)\n",
        "          cur_values = fwd_bx2[:,-i-1:-i if i > 0 else None]\n",
        "          cur_values = torch.repeat_interleave(cur_values, self.bkw_beam_width**i, 0)\n",
        "          cur_scores = fwd_scores[:,-i-1:-i if i > 0 else None]\n",
        "          cur_scores = torch.repeat_interleave(cur_scores, self.bkw_beam_width**i, 0)\n",
        "          cur_scores[mask] = float('-inf')  # mask from prev step\n",
        "          best_scores, best_beams = torch.topk(cur_scores, self.bkw_beam_width, 0)  # best beams\n",
        "          best_values = cur_values[best_beams]\n",
        "          mask = cur_values != best_values\n",
        "          mask = mask.view(new_batch_size, -1)\n",
        "          best_values = best_values.view(self.bkw_beam_width, -1)\n",
        "          best_values = torch.repeat_interleave(best_values, prev_batch_size, 0)\n",
        "          bkw_bx2 = bkw_bx2.repeat(self.bkw_beam_width, 1)  # increase batch for new beams\n",
        "          bkw_bx2 = torch.cat((best_values, bkw_bx2), 1)\n",
        "          best_scores = torch.repeat_interleave(best_scores, prev_batch_size, 0)\n",
        "          bkw_scores = bkw_scores.repeat(self.bkw_beam_width, 1)  # increase batch for new beams\n",
        "          bkw_scores = torch.cat((best_scores, bkw_scores), 1)\n",
        "        # bkw batch reduction\n",
        "        bkw_scores = self.depth_reduction(bkw_scores, axis=1) # cumulative beams scores\n",
        "        best_beams = torch.argmax(bkw_scores, axis=0, keepdim=True)  # best beams\n",
        "        bkw_bx2 = bkw_bx2[best_beams]  # best beam of all batches\n",
        "        bkw_bx2 = bkw_bx2.repeat(batch_size, 1)  # return to input batch size\n",
        "        bx2 = torch.cat((bx2, bkw_bx2), 1)  # <sos> + bkw_bx2\n",
        "    return bx2[:,1:], attn"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}