ria_rnn_bpe.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ria_rnn_bpe.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyP6Wg+c+jrqvLGa0cOVZayg",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "e4e5f0adb7b644baac8d38f789f63a57": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_33641e87cc1f4c9cbcbfb1b2a717e140",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_51565623338d46b6b2e50a1590e4020b",
              "IPY_MODEL_6711e9092a2b42edab6b4d5c62cecc0e"
            ]
          }
        },
        "33641e87cc1f4c9cbcbfb1b2a717e140": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "51565623338d46b6b2e50a1590e4020b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_a484f10715ed4ad7adc650e3b88efae4",
            "_dom_classes": [],
            "description": "loading samples: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "danger",
            "max": 1003869,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1003175,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_61c83d1d8bd44220a792d1c8ff696112"
          }
        },
        "6711e9092a2b42edab6b4d5c62cecc0e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_a11a298a3dd841bd9ed7032fd81ac356",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1003175/1003869 [02:10&lt;00:00, 10794.21it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_aff9bb01280b4cc3bb54d621c02d1440"
          }
        },
        "a484f10715ed4ad7adc650e3b88efae4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "61c83d1d8bd44220a792d1c8ff696112": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "a11a298a3dd841bd9ed7032fd81ac356": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "aff9bb01280b4cc3bb54d621c02d1440": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "e60ac13545ee4a17a94af84bbf022c30": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_98cab96334884deb9c3a97de6a92f213",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_3793b053272b46789944fae86bfeeffc",
              "IPY_MODEL_0e66f13265c249b6a6e2c32a75489077"
            ]
          }
        },
        "98cab96334884deb9c3a97de6a92f213": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "3793b053272b46789944fae86bfeeffc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_a42bd4c2eab348c4b9c220b928c58680",
            "_dom_classes": [],
            "description": "  0%",
            "_model_name": "FloatProgressModel",
            "bar_style": "danger",
            "max": 1,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 0,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_45070e65cb1745489fa9e8a206d0e54a"
          }
        },
        "0e66f13265c249b6a6e2c32a75489077": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_c1479450481e4a008b802dca3b353bd8",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0/1 [00:00&lt;?, ?it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_fefb8162039a424690605fb25d300797"
          }
        },
        "a42bd4c2eab348c4b9c220b928c58680": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "45070e65cb1745489fa9e8a206d0e54a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "c1479450481e4a008b802dca3b353bd8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "fefb8162039a424690605fb25d300797": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    },
    "accelerator": "GPU"
  },
  "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/ria_rnn_bpe.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FO83gr3P1Cqc",
        "colab_type": "text"
      },
      "source": [
        "# Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QgutxGKrctFP",
        "colab_type": "code",
        "outputId": "230f35a9-0ba4-423e-a420-b6ef0fcfd6a1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 178
        }
      },
      "source": [
        "from google.colab import drive\n",
        "import os\n",
        "\n",
        "drive.mount('/content/drive')\n",
        "root = '/content/drive/My Drive/'\n",
        "data_fn = 'datasets/norm_sents_ria.json.gz'\n",
        "vocab_fn = 'stop_lem_norm_sents_ria.json_vocab.json'\n",
        "enc_fn = 'datasets/norm_sents_ria.json_bpe_enc.json.gz'\n",
        "bpe_fn = 'bpe_ria.model'\n",
        "data_path = os.path.join(root, data_fn)\n",
        "vocab_path = os.path.join(root, vocab_fn)\n",
        "enc_path = os.path.join(root, enc_fn)\n",
        "bpe_path = os.path.join(root, bpe_fn)\n",
        "print(f'Check {data_path}: {os.path.exists(data_path)}')\n",
        "print(f'Check {vocab_path}: {os.path.exists(vocab_path)}')\n",
        "print(f'Check {enc_path}: {os.path.exists(enc_path)}')"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3aietf%3awg%3aoauth%3a2.0%3aoob&response_type=code&scope=email%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdocs.test%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive.photos.readonly%20https%3a%2f%2fwww.googleapis.com%2fauth%2fpeopleapi.readonly\n",
            "\n",
            "Enter your authorization code:\n",
            "··········\n",
            "Mounted at /content/drive\n",
            "Check /content/drive/My Drive/datasets/norm_sents_ria.json.gz: True\n",
            "Check /content/drive/My Drive/stop_lem_norm_sents_ria.json_vocab.json: True\n",
            "Check /content/drive/My Drive/datasets/norm_sents_ria.json_bpe_enc.json.gz: True\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GGlUFUDRqMiI",
        "colab_type": "text"
      },
      "source": [
        "## BPE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oo-lnkHKtdkV",
        "colab_type": "code",
        "outputId": "7da3c1f4-dd04-4a74-cbd0-734678374b77",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 143
        }
      },
      "source": [
        "!pip install youtokentome"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting youtokentome\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/a3/65/4a86cf99da3f680497ae132329025b291e2fda22327e8da6a9476e51acb1/youtokentome-1.0.6-cp36-cp36m-manylinux2010_x86_64.whl (1.7MB)\n",
            "\u001b[K     |████████████████████████████████| 1.7MB 6.3MB/s \n",
            "\u001b[?25hRequirement already satisfied: Click>=7.0 in /usr/local/lib/python3.6/dist-packages (from youtokentome) (7.1.2)\n",
            "Installing collected packages: youtokentome\n",
            "Successfully installed youtokentome-1.0.6\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LzMkIQi1CeFh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# import gzip \n",
        "# import json\n",
        "# from tqdm.notebook import tqdm\n",
        "\n",
        "# bpe_train_text = root + 'bpe_train_text.txt'\n",
        "# with gzip.open(root + 'datasets/norm_sents_ria.json.gz', 'rb') as f:\n",
        "#   with open(bpe_train_text, 'w') as f2:\n",
        "#     for line in tqdm(f):\n",
        "#       n = json.loads(line)\n",
        "#       for s in n['text']:\n",
        "#         f2.write(s + '\\n')\n",
        "#       for s in n['title']:\n",
        "#         f2.write(s + '\\n')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s97dRzCpthq4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# import random\n",
        "# import youtokentome as yttm\n",
        "\n",
        "# random.seed(0)\n",
        "\n",
        "# yttm.BPE.train(data=bpe_train_text, model=bpe_path,\n",
        "#               vocab_size=50000, coverage=0.9999, n_threads=-1, \n",
        "#               pad_id=0, unk_id=1, bos_id=2, eos_id=3)\n",
        "\n",
        "# # Loading model\n",
        "# bpe = yttm.BPE(model=bpe_path)\n",
        "\n",
        "# print(bpe.vocab())\n",
        "# print(test_text)\n",
        "# # Two types of tokenization\n",
        "# print(bpe.encode([test_text], output_type=yttm.OutputType.ID, bos=True, eos=True))\n",
        "# enc_text = bpe.encode([test_text], output_type=yttm.OutputType.ID)\n",
        "# print(bpe.encode([test_text], output_type=yttm.OutputType.SUBWORD))\n",
        "# print(bpe.decode(enc_text, ignore_ids=None))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xdzcMNI1qRD3",
        "colab_type": "text"
      },
      "source": [
        "## DS"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dRlGjfpXdU88",
        "colab_type": "code",
        "outputId": "fd87abf3-2d06-4e50-b1c0-1ec43775ca12",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 157,
          "referenced_widgets": [
            "e4e5f0adb7b644baac8d38f789f63a57",
            "33641e87cc1f4c9cbcbfb1b2a717e140",
            "51565623338d46b6b2e50a1590e4020b",
            "6711e9092a2b42edab6b4d5c62cecc0e",
            "a484f10715ed4ad7adc650e3b88efae4",
            "61c83d1d8bd44220a792d1c8ff696112",
            "a11a298a3dd841bd9ed7032fd81ac356",
            "aff9bb01280b4cc3bb54d621c02d1440"
          ]
        }
      },
      "source": [
        "import torch\n",
        "from torch.utils.data import Dataset, DataLoader, random_split\n",
        "import os\n",
        "import json\n",
        "import gzip\n",
        "from pprint import pprint\n",
        "from tqdm.notebook import tqdm\n",
        "import numpy as np\n",
        "import gc\n",
        "import shutil\n",
        "import youtokentome as yttm\n",
        "\n",
        "class RiaDataset(Dataset):\n",
        "  def __init__(self, data=None, max_vocab=10000, max_samples=500,\n",
        "               vocab=None, enc_data=None, bpe=None,\n",
        "               special_chars=('<pad>', '<unk>', '<sos>', '<eos>')):\n",
        "    super(RiaDataset).__init__()\n",
        "    if isinstance(bpe, str):  # Load vocabualry\n",
        "      self.bpe = yttm.BPE(model=bpe)\n",
        "    elif bpe is not None:\n",
        "      self.bpe = bpe\n",
        "    # Vocabulary dict {word: count)\n",
        "    if vocab is None and bpe is None:  # Create vocabualry\n",
        "      print('Creating vocabulary')\n",
        "      if data is None:\n",
        "        raise ValueError('Requires raw data path for vocabulary creating')\n",
        "      vocab = {}\n",
        "      cnt, bad_ids = 0, []\n",
        "      root, ext = os.path.splitext(data)\n",
        "      with gzip.open(data, 'rb') if ext == '.gz' else open(data) as f:\n",
        "        try:\n",
        "          for _ in tqdm(range(max_samples), desc='loading samples'):\n",
        "            n = json.loads(next(f))\n",
        "            if not len(n['text']) == 0 and not len(n['title']) == 0:\n",
        "              for s in n['text']:\n",
        "                for w in s.split():\n",
        "                  vocab[w] = vocab.setdefault(w, 0) + 1\n",
        "              for s in n['title']:\n",
        "                for w in s.split():\n",
        "                  vocab[w] = vocab.setdefault(w, 0) + 1\n",
        "            else:\n",
        "              bad_ids.append(cnt)\n",
        "            cnt += 1\n",
        "        except StopIteration as e:\n",
        "          print(f'max_samples {max_samples} > len dataset {cnt}')\n",
        "      print(f'bad texts: {len(bad_ids)} {bad_ids}')\n",
        "      vocab_path = root + '_vocab.json'\n",
        "      print(f'Saving vocabulary to {vocab_path}')\n",
        "      with open(vocab_path, 'w') as f:\n",
        "        json.dump(vocab, f)\n",
        "    elif isinstance(vocab, str):  # Load vocabualry\n",
        "      print(f'Loading vocabulary from {vocab}')\n",
        "      with open(vocab) as f:\n",
        "        vocab = json.load(f)\n",
        "    elif bpe is None:   # else use as it is   \n",
        "      print('Full vocabulary')\n",
        "      print(f'unique: {len(vocab)}, total: {sum(vocab.values())}')\n",
        "      self.itos = list(special_chars)\n",
        "      top_words = [(w, c) for w, c in sorted(vocab.items(), key=lambda x: x[1],\n",
        "                                              reverse=True)][:max_vocab]\n",
        "      self.itos += [w for w, _ in top_words]\n",
        "      self.stoi = {x: i for i, x in enumerate(self.itos)}\n",
        "      assert len(self.itos) == len(self.stoi)\n",
        "      print('Reduced vocabulary')\n",
        "      print(f'unique: {len(self.stoi)}, total: {sum(c for _, c in top_words)} ~ \\\n",
        "      {sum(c for _, c in top_words)/sum(vocab.values())*100:6.3f}% of full')\n",
        "    # Samples [list of np.arrays (text), list of np.arrays (title)]\n",
        "    self.samples = [] \n",
        "    if enc_data is None:  # Encode samples\n",
        "      print('Encoding samples')\n",
        "      if data is None:\n",
        "        raise ValueError('Requires raw data path for samples encoding')\n",
        "      cnt, bad_ids = 0, []\n",
        "      root, ext = os.path.splitext(data)\n",
        "      with gzip.open(data, 'rb') if ext == '.gz' else open(data) as f:\n",
        "        try:\n",
        "          for _ in tqdm(range(max_samples), desc='loading samples'):\n",
        "            n = json.loads(next(f))\n",
        "            if not len(n['text']) == 0 and not len(n['title']) == 0:        \n",
        "              x1 = [self.encode(x) for x in n['text']]\n",
        "              x2 = [self.encode(x) for x in n['title']]\n",
        "              self.samples.append([x1, x2])\n",
        "            else:\n",
        "              bad_ids.append(cnt)\n",
        "            cnt += 1\n",
        "        except StopIteration as e:\n",
        "            print(f'max_samples {max_samples} > len dataset {cnt}')\n",
        "      print(f'bad texts: {len(bad_ids)} {bad_ids}')\n",
        "      enc_path = root + '_bpe_enc.json'\n",
        "      print(f'Saving encoded samples to {enc_path}')\n",
        "      with open(enc_path, 'w') as f:\n",
        "        for s in tqdm(self.samples):\n",
        "          json_str = json.dumps(s, default=lambda x: x.tolist())\n",
        "          f.write(json_str + '\\n')\n",
        "      comp_enc_path = enc_path + '.gz'\n",
        "      print(f'Compressing {enc_path} to {comp_enc_path}')\n",
        "      with gzip.open(comp_enc_path, 'wb') as gz_file:\n",
        "        with open(enc_path, 'rb') as json_file:\n",
        "          shutil.copyfileobj(json_file, gz_file)\n",
        "      print(f'Deleting uncompressed samples: {enc_path}')\n",
        "      os.remove(enc_path)\n",
        "    elif isinstance(enc_data, str):  # Load encodings\n",
        "      print(f'Loading encoded samples from {enc_data}')\n",
        "      _, ext = os.path.splitext(enc_data)\n",
        "      with gzip.open(enc_data, 'rb') if ext == '.gz' else open(enc_data) as f:\n",
        "        try:\n",
        "          for _ in tqdm(range(max_samples), desc='loading samples'):\n",
        "            s = json.loads(next(f))\n",
        "            x1 = [np.array(x) for x in s[0]]\n",
        "            x2 = [np.array(x) for x in s[1]]\n",
        "            self.samples.append([x1, x2])\n",
        "        except StopIteration as e:\n",
        "          print(f'max_samples {max_samples} > len dataset {len(self.samples)}')\n",
        "    else:\n",
        "      self.samples = enc_data\n",
        "    print(f'samples: {len(self.samples)}')\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.samples)\n",
        "\n",
        "  def __getitem__(self, i):\n",
        "    return self.samples[i]\n",
        "\n",
        "  def encode(self, s):\n",
        "    # return np.array([self.stoi['<sos>']] \\\n",
        "    #               + [self.stoi.get(x, self.stoi['<unk>']) for x in s.split()] \\\n",
        "    #               + [self.stoi['<eos>']])\n",
        "    return np.array(self.bpe.encode(s, output_type=yttm.OutputType.ID,\n",
        "                                    bos=True, eos=True))\n",
        "  def decode(self, s):\n",
        "    # return [self.itos[x] for x in s]\n",
        "    return self.bpe.decode(s.tolist() if not isinstance(s, list) else s,\n",
        "                           ignore_ids=None)\n",
        "\n",
        "class Collate():\n",
        "  def __init__(self, n_x1=1, padding_value=0):\n",
        "     self.n_x1 = n_x1\n",
        "     self.padding_value = padding_value\n",
        "\n",
        "  def __call__(self, batch):\n",
        "    bx1, bx2 = [], []\n",
        "    for b in batch:\n",
        "      x1s, x2s = b\n",
        "      # skip first sentence (usually date and place)\n",
        "      ei = min(len(x1s), 1 + self.n_x1)\n",
        "      si = 1 if ei > 1 else 0\n",
        "      # print(self.n_x1, si, ei, len(x1s))\n",
        "      for x1 in x1s[si:ei]:\n",
        "        for x2 in x2s:\n",
        "          bx1.append(torch.as_tensor(x1))\n",
        "          bx2.append(torch.as_tensor(x2))\n",
        "    bx1 = torch.nn.utils.rnn.pad_sequence(bx1, batch_first=True, \n",
        "                                          padding_value=0)\n",
        "    bx2 = torch.nn.utils.rnn.pad_sequence(bx2, batch_first=True,\n",
        "                                          padding_value=0)\n",
        "    batch = [bx1, bx2]\n",
        "    return batch\n",
        "\n",
        "# Create vocabulary and encodings  # max_samples=1003869 max_vocab=50000\n",
        "# ds = RiaDataset(data=data_path, max_vocab=50000, max_samples=1003869, \n",
        "#                 vocab=None, enc_data=None, bpe=bpe_path)\n",
        "ds = RiaDataset(data=data_path, max_vocab=50000, max_samples=1003869, \n",
        "                vocab=None, enc_data=enc_path, bpe=bpe_path)\n",
        "# ds = RiaDataset(data=data_path, max_vocab=50000, max_samples=2003869, \n",
        "#                 vocab=None, enc_data=None)\n",
        "# Load vocabulary and encodings\n",
        "# ds = RiaDataset(data=None, max_vocab=50000, max_samples=1003869, \n",
        "#                 bpe=bpe_path,\n",
        "#                 vocab=vocab_path, enc_data=enc_path) \n",
        "# unique words: 601956, total words: 198199252\n",
        "# vocabulary: 50004, words in vocabulary: 194385367\n",
        "\n",
        "torch.manual_seed(0)\n",
        "torch.cuda.manual_seed(0)\n",
        "# dl = DataLoader(ds, batch_size=1, num_workers=1, shuffle=False, \n",
        "#                 drop_last=False, collate_fn=Collate())\n",
        "# for s in tqdm(dl, desc='dataset'):\n",
        "#   bx1, bx2 = s\n",
        "#   assert bx1.size()[0] == bx2.size()[0]\n",
        "#   print(bx1.size(), bx2.size())\n",
        "#   for x1, x2 in zip(bx1, bx2):\n",
        "#     print(ds.decode(x1), ds.decode(x2))\n",
        "#   print([ds.decode(x) for x in bx1])\n",
        "#   print([ds.decode(x) for x in bx2])\n",
        "train_len = int(0.7*len(ds))\n",
        "test_len = int(0.2*len(ds))\n",
        "val_len = len(ds) - train_len - test_len\n",
        "lens = [train_len, test_len, val_len]\n",
        "print(lens, sum(lens), len(ds))\n",
        "train_ds, test_ds, val_ds = random_split(ds, lens)\n",
        "train_dl = DataLoader(train_ds, batch_size=2, num_workers=1, \n",
        "                      shuffle=True, drop_last=False, \n",
        "                      collate_fn=Collate(n_x1=1))\n",
        "test_dl = DataLoader(test_ds, batch_size=3, num_workers=1, \n",
        "                     shuffle=False, drop_last=False, \n",
        "                     collate_fn=Collate(n_x1=1))\n",
        "val_dl = DataLoader(val_ds, batch_size=3, num_workers=1, \n",
        "                    shuffle=False, drop_last=False, \n",
        "                    collate_fn=Collate(n_x1=1))\n",
        "# for s in tqdm(train_dl, desc='train'):\n",
        "#   bx1, bx2 = s\n",
        "#   assert bx1.size()[0] == bx2.size()[0]\n",
        "#   # print(bx1.size(), bx2.size())\n",
        "# for s in tqdm(test_dl, desc='test'):\n",
        "#   bx1, bx2 = s\n",
        "#   assert bx1.size()[0] == bx2.size()[0]\n",
        "#   # print(bx1.size(), bx2.size())\n",
        "# for s in tqdm(val_dl, desc='val'):\n",
        "#   bx1, bx2 = s\n",
        "#   assert bx1.size()[0] == bx2.size()[0]\n",
        "#   # print(bx1.size(), bx2.size())\n",
        "gc.collect()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loading encoded samples from /content/drive/My Drive/datasets/norm_sents_ria.json_bpe_enc.json.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "e4e5f0adb7b644baac8d38f789f63a57",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='loading samples', max=1003869.0, style=ProgressStyle(desc…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "max_samples 1003869 > len dataset 1003723\n",
            "samples: 1003723\n",
            "[702606, 200744, 100373] 1003723 1003723\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eYgtjOABt2rJ",
        "colab_type": "code",
        "outputId": "4766c727-6c37-4726-9670-24d9d517ceda",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "# import matplotlib.pyplot as plt\n",
        "# import seaborn as sns\n",
        "# import pandas as pd\n",
        "\n",
        "# cnt = sum(len(y) for x in ds.samples for y in x[0]) + sum(len(y) for x in ds.samples for y in x[1])\n",
        "# print(cnt)\n",
        "# # max_vocab = 1000\n",
        "# # with open(vocab_path) as f:\n",
        "# #   vocab = json.load(f)\n",
        "# # top_words = dict([(w, c) for w, c in sorted(vocab.items(), key=lambda x: x[1],\n",
        "# #                                             reverse=True)][:max_vocab])\n",
        "# # # # sns.distplot(list(map(len, vocab.keys())), kde=False)\n",
        "# # # # plt.figure()\n",
        "# # # # sns.distplot(list(map(len, top_words.keys())), kde=True)\n",
        "# # # # plt.figure()\n",
        "# # sns.distplot(list(vocab.values()), kde=False, hist_kws={'log': True})\n",
        "# # sns.distplot(list(top_words.values()), kde=False, hist_kws={'log': True})\n",
        "# # ax = sns.distplot(list(vocab.values()), kde=False, hist=True, \n",
        "# #              hist_kws={'log': True}, norm_hist=False)\n",
        "# # # ax.set_xscale('log')\n",
        "# # ax = sns.distplot(list(top_words.values()), kde=False, hist=True, \n",
        "# #              hist_kws={'log': True}, norm_hist=False)\n",
        "# # # ax.set_xscale('log')\n",
        "# # plt.figure()\n",
        "# # plt.title('Sentences lengths')\n",
        "# # sns.kdeplot([len(y) for x in ds.samples for y in x[0]], \n",
        "# #             label='texts', shade=True, bw=2)\n",
        "# # sns.kdeplot([len(y) for x in ds.samples for y in x[1]], \n",
        "# #             label='titles', shade=True, bw=2)\n",
        "# # plt.legend()\n",
        "# # plt.figure()\n",
        "\n",
        "\n",
        "\n",
        "# # sns.catplot(data=df, orient=\"h\", kind=\"swarm\")\n",
        "\n",
        "# # plt.figure()\n",
        "# # sns.distplot([len(x[0][0]) for x in ds.samples], norm_hist=False, \n",
        "# #              hist_kws={'log': True}, \n",
        "# #              kde=False)\n",
        "# # plt.figure()\n",
        "# # sns.distplot([len(x[1][0]) for x in ds.samples], \n",
        "# #             #  hist_kws={'log': True}, \n",
        "# #              kde=True, norm_hist=False)\n",
        "\n",
        "# df = data=pd.DataFrame({\n",
        "#     \"texts 1 sentense length\": [len(x[0][1]) if len(x[0]) > 1 else len(x[0][0]) \n",
        "#     for x in ds.samples[:100000]],\n",
        "#     \"titles 1 sentense length\": [len(x[1][0]) for x in ds.samples[:100000]]}\n",
        "# )\n",
        "# # sns.jointplot(y='titles 1 sentense length', x='texts 1 sentense length', \n",
        "# #               data=df, kind='kde', bw=2.)\n",
        "# sns.jointplot(y='titles 1 sentense length', x='texts 1 sentense length', \n",
        "#               data=df)\n",
        "# # plt.figure()\n",
        "# # sns.kdeplot([len(x[0][1]) if len(x[0]) > 1 else len(x[0][0])\n",
        "# #             for x in ds.samples], bw=5., shade=True)\n",
        "# # sns.kdeplot([len(x[1][0]) for x in ds.samples], bw=5., shade=True)\n",
        "\n",
        "\n",
        "# # ax = sns.kdeplot([len(x[0][1]) if len(x[0]) > 1 else len(x[0][0]) \n",
        "# #             for x in ds.samples[:100000]], \n",
        "# #             [len(x[1][0]) for x in ds.samples[:100000]],\n",
        "# #             n_levels=3, cmap='Greys_r',\n",
        "# #             bw=2., shade=True, clip=(0, 1000), shade_lowest=True)\n",
        "# # ax.set_xlabel('text first sentence length')\n",
        "# # ax.set_ylabel('title first sentence length')\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "330582350\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bwYBLWZRXQqS",
        "colab_type": "text"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AfI1NnKX9c1o",
        "colab_type": "code",
        "cellView": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n",
        "\n",
        "class EncoderRNN(nn.Module):\n",
        "  def __init__(self, num_embeddings, embedding_dim, padding_idx,\n",
        "               rnn_type, hidden_size, num_layers=1, rnn_dropout=0,\n",
        "               bidirectional=False, dropout=0, pack=False):\n",
        "    super(EncoderRNN, self).__init__()\n",
        "    rnn_map = {'RNN': nn.RNN, 'GRU': nn.GRU, 'LSTM': nn.LSTM}\n",
        "    self.embedding = nn.Embedding(num_embeddings=num_embeddings, \n",
        "                                  embedding_dim=embedding_dim,\n",
        "                                  padding_idx=padding_idx)\n",
        "    self.rnn = rnn_map[rnn_type](input_size=embedding_dim, \n",
        "                                 hidden_size=hidden_size,\n",
        "                                 num_layers=num_layers,\n",
        "                                 batch_first=True,\n",
        "                                 dropout=rnn_dropout,\n",
        "                                 bidirectional=bidirectional)\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "    self.pack = pack\n",
        "\n",
        "  def forward(self, x1):  # [B, L]\n",
        "    x = self.embedding(x1)  # [B, L] -> [B, L, E]\n",
        "    if self.pack:  # [B, L, E] -> Packed [B, L, E]\n",
        "      ls = torch.sum(x1 != 0, dim=1)\n",
        "      if torch.any(ls != ls[0]):  # all batches has equal lengths\n",
        "        x = pack_padded_sequence(x, ls, batch_first=True, enforce_sorted=False)\n",
        "    # [B, L, E] -> (Packed) [B, L, ND*H], ([NL*ND, B, H], ([NL*ND, B, H]))\n",
        "    if isinstance(self.rnn, nn.LSTM):\n",
        "      ht, (hn, cn) = self.rnn(x)\n",
        "      hn, cn = self.dropout(hn), self.dropout(cn)\n",
        "      return ht, (hn, cn)\n",
        "    else:\n",
        "      ht, hn = self.rnn(x)\n",
        "      hn = self.dropout(hn)\n",
        "      return ht, hn\n",
        "\n",
        "\n",
        "class AttentionDecoderRNN(nn.Module):\n",
        "  def __init__(self, num_embeddings, embedding_dim, padding_idx,\n",
        "               rnn_type, hidden_size, num_layers=1, rnn_dropout=0,\n",
        "               bidirectional=False, dropout=0, out_hidden=0,\n",
        "               attn_type='soft_dot', pack=False):\n",
        "    super(AttentionDecoderRNN, self).__init__()\n",
        "    rnn_map = {'RNN': nn.RNN, 'GRU': nn.GRU, 'LSTM': nn.LSTM}\n",
        "    attn_types = ['dot', 'cos', 'dist', 'soft_dot', 'soft_cos', 'soft_dist', 'none']\n",
        "    self.attn_type = attn_type\n",
        "    self.embedding = nn.Embedding(num_embeddings=num_embeddings, \n",
        "                                  embedding_dim=embedding_dim,\n",
        "                                  padding_idx=padding_idx)\n",
        "    self.rnn = rnn_map[rnn_type](input_size=embedding_dim, \n",
        "                                 hidden_size=hidden_size,\n",
        "                                 num_layers=num_layers,\n",
        "                                 batch_first=True,\n",
        "                                 dropout=rnn_dropout,\n",
        "                                 bidirectional=bidirectional)\n",
        "    out_input = hidden_size if self.attn_type == 'none' else 2*hidden_size\n",
        "    if out_hidden > 0:\n",
        "      self.out_hidden = nn.Linear(out_input, out_hidden)\n",
        "      self.out = nn.Linear(out_hidden, num_embeddings)\n",
        "    else:\n",
        "      self.out_hidden = None\n",
        "      self.out = nn.Linear(out_input, num_embeddings)\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "    self.softmax = nn.LogSoftmax(dim=2)\n",
        "    self.pack = pack\n",
        "\n",
        "  def forward(self, x2, h1, x1):  # [B, L], [B, L, ND*H], ([NL*ND, B, H], [NL*ND, B, H])\n",
        "    x = self.embedding(x2)  # [B, L] -> [B, L, E]\n",
        "    if self.pack:  # [B, L, E] -> Packed [B, L, E]\n",
        "      ls2 = torch.sum(x2 != 0, dim=1)\n",
        "      if torch.any(ls2 != ls2[0]):\n",
        "        x = pack_padded_sequence(x, ls2, batch_first=True, enforce_sorted=False)\n",
        "    ht1, hns1 = h1  # [B, L, ND*H], ([NL*ND, B, H], [NL*ND, B, H])\n",
        "    # [B, L, D], ([NL*ND, B, H], [NL*ND, B, H]) -> [B, L, ND*H], ([NL*ND, B, H], [NL*ND, B, H])\n",
        "    ht2, hns2 = self.rnn(x, hns1)\n",
        "    if self.pack:  # Packed [B, L, E] -> [B, L, E]\n",
        "      ls1 = torch.sum(x1 != 0, dim=1)\n",
        "      if torch.any(ls1 != ls1[0]):  # all batches has equal lengths\n",
        "        ht1, _ = pad_packed_sequence(ht1, batch_first=True)\n",
        "      if torch.any(ls2 != ls2[0]):  # all batches has equal lengths\n",
        "        ht2, _ = pad_packed_sequence(ht2, batch_first=True)\n",
        "    if self.attn_type == 'none':\n",
        "      if self.rnn.bidirectional:\n",
        "        B2, L2, NDH2 = ht2.size()\n",
        "        # [B, L, ND*H] -> [B, L, ND, H] -> [B, L, H]\n",
        "        ht2 = ht2.view(B2, L2, 2, int(NDH2/2)).mean(2)\n",
        "      x = self.dropout(ht2)\n",
        "      if self.out_hidden is not None:\n",
        "        x = self.out_hidden(x)\n",
        "      x = self.out(x)\n",
        "      x = self.softmax(x)\n",
        "      return x, None\n",
        "    else:\n",
        "      if self.rnn.bidirectional:\n",
        "        B1, L1, NDH1 = ht1.size()\n",
        "        B2, L2, NDH2 = ht2.size()  # B1 == B2, NDH1 == NDH2\n",
        "        # [B, L, ND*H] -> [B, L, ND, H] -> [B, L, H]\n",
        "        ht1 = ht1.view(B1, L1, 2, int(NDH1/2)).mean(2)\n",
        "        ht2 = ht2.view(B2, L2, 2, int(NDH2/2)).mean(2)\n",
        "      # [L2, H], [L1, H] -> [L2, L1]\n",
        "      # mask where x1 and x2 token is <PAD>\n",
        "      # pad_mask = torch.einsum('bi,bj->bij', x2, x1) == 0\n",
        "      # mask where only x1 token is <PAD>\n",
        "      pad_mask = torch.einsum('bi,bj->bij', torch.ones_like(x2), x1) == 0\n",
        "      # [B, L2, H], [B, L1, H] -> [B, L2, L1]  # attention\n",
        "      # [B, L2, L1], [B, L1, H] -> [B, L2, H]  # weighted h1\n",
        "      if self.attn_type == 'dot':\n",
        "        attn = torch.einsum('bih,bjh->bij', ht2, ht1)  # dot product\n",
        "        attn[pad_mask] = 0\n",
        "      elif self.attn_type == 'soft_dot':\n",
        "        attn = torch.einsum('bih,bjh->bij', ht2, ht1)  # dot product\n",
        "        attn[pad_mask] = float('-inf')\n",
        "        attn = F.softmax(attn, 2)\n",
        "      elif self.attn_type == 'dist':\n",
        "        ht1, ht2 = ht1.contiguous(), ht2.contiguous()\n",
        "        attn = torch.cdist(ht2, ht1)  # euclidian distance\n",
        "        attn = torch.masked_fill(attn, pad_mask, float('inf'))\n",
        "        attn = F.threshold(attn, threshold=1e-6, value=1e-6)  # short dist\n",
        "        attn = 1/attn  # inverse\n",
        "        attn = F.normalize(attn, p=1, dim=2)  # to [0, 1]\n",
        "      elif self.attn_type == 'soft_dist':\n",
        "        ht1, ht2 = ht1.contiguous(), ht2.contiguous()\n",
        "        attn = torch.cdist(ht2, ht1)  # euclidian distance\n",
        "        attn = torch.masked_fill(attn, pad_mask, float('inf'))\n",
        "        attn = F.softmin(attn, 2)\n",
        "      elif self.attn_type == 'cos':\n",
        "        ht1n = F.normalize(ht1, p=2, dim=2)  # normalize to length = 1\n",
        "        ht2n = F.normalize(ht2, p=2, dim=2)  # normalize to length = 1\n",
        "        attn = torch.einsum('bih,bjh->bij', ht2n, ht1n)  # dot product\n",
        "        attn[pad_mask] = 0\n",
        "      elif self.attn_type == 'soft_cos':\n",
        "        ht1n = F.normalize(ht1, p=2, dim=2)  # normalize to length = 1\n",
        "        ht2n = F.normalize(ht2, p=2, dim=2)  # normalize to length = 1\n",
        "        attn = torch.einsum('bih,bjh->bij', ht2n, ht1n)  # dot product\n",
        "        attn[pad_mask] = float('-inf')\n",
        "        attn = F.softmax(attn, 2)\n",
        "      hw1 = torch.einsum('bij,bjh->bih', attn, ht1)  # weighted h1\n",
        "      ha = torch.cat((ht2, hw1), 2)  # [B, L2, H], [B, L2, H] -> [B, L2, H+H]\n",
        "      x = self.dropout(ha)  # [B, L2, H+H] -> [B, L2, H+H]\n",
        "      if self.out_hidden is not None:\n",
        "        x = self.out_hidden(x)   # [B, L2, H+H] -> [B, L2, OH]\n",
        "      x = self.out(x)  # [B, L2, H+H] or [B, L2, OH] -> [B, L2, D2]\n",
        "      y2 = self.softmax(x)  # [B, L2, D2] -> [B, L2, D2]\n",
        "      return y2, attn\n",
        "\n",
        "\n",
        "class EncoderDecoder(nn.Module):\n",
        "  def __init__(self, \n",
        "               enc_num_embeddings, enc_embedding_dim, enc_padding_idx,\n",
        "               dec_num_embeddings, dec_embedding_dim, dec_padding_idx,\n",
        "               rnn_type, hidden_size, num_layers=1, out_hidden=0,\n",
        "               enc_rnn_dropout=0, dec_rnn_dropout=0,\n",
        "               bidirectional=False, enc_dropout=0, dec_dropout=0, \n",
        "               attn_type='dot', pack=False):\n",
        "    super(EncoderDecoder, self).__init__()\n",
        "    self.encoder = EncoderRNN(num_embeddings=enc_num_embeddings, \n",
        "                              embedding_dim=enc_embedding_dim, \n",
        "                              padding_idx=enc_padding_idx,\n",
        "                              hidden_size=hidden_size, \n",
        "                              rnn_type=rnn_type, \n",
        "                              bidirectional=bidirectional,\n",
        "                              num_layers=num_layers,\n",
        "                              dropout=enc_dropout,\n",
        "                              rnn_dropout=enc_rnn_dropout,\n",
        "                              pack=pack)\n",
        "    self.decoder = AttentionDecoderRNN(num_embeddings=dec_num_embeddings, \n",
        "                                       embedding_dim=dec_embedding_dim, \n",
        "                                       padding_idx=dec_padding_idx,\n",
        "                                       hidden_size=hidden_size, \n",
        "                                       rnn_type=rnn_type, \n",
        "                                       bidirectional=bidirectional,\n",
        "                                       num_layers=num_layers,\n",
        "                                       dropout=dec_dropout,\n",
        "                                       out_hidden=out_hidden,\n",
        "                                       rnn_dropout=dec_rnn_dropout,\n",
        "                                       attn_type=attn_type,\n",
        "                                       pack=pack)\n",
        "\n",
        "  def forward(self, x1, x2):\n",
        "    #  [B, L1] -> [B, L1, ND*H], ([NL*ND, B, H], [NL*ND, B, H])\n",
        "    h1 = self.encoder(x1)\n",
        "    # [B, L2], ([B, L1, ND*H], ([NL*ND, B, H], [NL*ND, B, H])) -> [B, L2, E2]\n",
        "    y2 = self.decoder(x2, h1, x1)\n",
        "    return y2\n",
        "\n",
        "\n",
        "def external_attn(x1, x2, attn_dict):\n",
        "    attn_dict = {1: [1], 2: [2], 4: [4, 3]}\n",
        "    attn = []\n",
        "    B1, L1 = x1.size()\n",
        "    B2, L2 = x2.size()\n",
        "    for i in range(B1):\n",
        "      b = []\n",
        "      for j in range(L2):\n",
        "        l2 = []\n",
        "        # x2t = x2[i, j].item()\n",
        "        x2t = x2[i, j + 1].item() if j + 2 < L2 else 0 # decoder shift\n",
        "        for k in range(L1):\n",
        "          x1t = x1[i, k].item()\n",
        "          # x1t = x1[i, k + 1].item() if k + 1 < L1 else 0 # encoder shift\n",
        "          x1ts = self.attn_dict.get(x2t, [])\n",
        "          if x1t in x1ts:\n",
        "            l2.append(1.)\n",
        "          else:\n",
        "            l2.append(0.)\n",
        "        b.append(l2)\n",
        "      attn.append(b)\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    attn = torch.tensor(attn).to(device)\n",
        "    attn = F.normalize(attn, p=1, dim=2)  # to [0, 1]\n",
        "    return attn\n",
        "\n",
        "from IPython.display import Image\n",
        "# Image(make_dot(loss).render('loss', format='png'))\n",
        "from tqdm.notebook import tqdm\n",
        "from torch import optim\n",
        "import numpy as np\n",
        "\n",
        "def plot_attention(a, x1, x2, shift=True, mask=True, \n",
        "                   suptitle='Attention', figsize=None, \n",
        "                   decoder_x1=None, decoder_x2=None, tight=False,\n",
        "                   labels=True):\n",
        "  # %matplotlib inline\n",
        "  import matplotlib.pyplot as plt\n",
        "  from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
        "\n",
        "  b = a.shape[0]\n",
        "  fig, axs = plt.subplots(1, b, figsize=figsize)\n",
        "  if not isinstance(axs, np.ndarray):  # if batch size == 1\n",
        "    axs = [axs]\n",
        "  fig.suptitle(suptitle)\n",
        "  for i in range(b):\n",
        "    ax = axs[i]\n",
        "    if shift:\n",
        "      ba, bx1, bx2 = a[i,:-1,1:], x1[i,1:], x2[i,1:]\n",
        "    else:\n",
        "      ba, bx1, bx2 = a[i], x1[i], x2[i]\n",
        "    if mask:\n",
        "      mask_x1 = np.flatnonzero(bx1)\n",
        "      mask_x2 = np.flatnonzero(bx2)\n",
        "      ba, bx1, bx2 =  ba[mask_x2,:][:,mask_x1], bx1[mask_x1], bx2[mask_x2]\n",
        "    # ax.set_title(f'{i+1}', y=-0.2)\n",
        "    im = ax.imshow(ba, cmap='gray')\n",
        "    ax.set_xticks(np.arange(len(bx1)))\n",
        "    ax.set_yticks(np.arange(len(bx2)))\n",
        "    if labels:\n",
        "      ax.set_xlabel('x1')\n",
        "      ax.set_ylabel('x2')\n",
        "    ax.xaxis.tick_top()\n",
        "    ax.xaxis.set_label_position('top')\n",
        "    if decoder_x1 is not None:\n",
        "      bx1 = decoder_x1(bx1)\n",
        "    if decoder_x2 is not None:\n",
        "      bx2 = decoder_x2(bx2)\n",
        "    ax.set_xticklabels(bx1, rotation=90)  # rotation=90\n",
        "    ax.set_yticklabels(bx2)\n",
        "    fig.colorbar(im, ax=ax, fraction=0.05, pad=0.05)\n",
        "  if tight:\n",
        "    plt.tight_layout()\n",
        "\n",
        "\n",
        "# torch.backends.cudnn.enabled=False\n",
        "# torch.backends.cudnn.deterministic=True\n",
        "# torch.autograd.set_detect_anomaly(True)\n",
        "\n",
        "# enc_dec_config = {\n",
        "#   'enc_num_embeddings': 6,\n",
        "#   'enc_embedding_dim': 2,\n",
        "#   'enc_padding_idx': 0,\n",
        "#   'dec_num_embeddings': 6,\n",
        "#   'dec_embedding_dim': 2,\n",
        "#   'dec_padding_idx': 0,\n",
        "#   'rnn_type': 'RNN',\n",
        "#   'hidden_size': 2,\n",
        "#   'num_layers': 1,\n",
        "#   'bidirectional': False,\n",
        "#   'enc_rnn_dropout': 0,\n",
        "#   'dec_rnn_dropout': 0,\n",
        "#   'enc_dropout': 0,\n",
        "#   'dec_dropout': 0,\n",
        "#   'attn_type': 'soft_dist',\n",
        "#   'out_hidden': 16,\n",
        "#   'pack': True\n",
        "# }\n",
        "# # <pad>, <unk>, <go>, <eos>\n",
        "# device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "# seed = 0\n",
        "# torch.manual_seed(seed)\n",
        "# torch.cuda.manual_seed(seed)\n",
        "# model = EncoderDecoder(**enc_dec_config).to(device)\n",
        "# x1 = torch.tensor([[2, 1, 4, 3, 0], [2, 1, 4, 1, 3]]).to(device)\n",
        "# x2 = torch.tensor([[2, 4, 1, 5, 3, 0], [2, 4, 5, 1, 1, 3]]).to(device)\n",
        "# opt = optim.Rprop(model.parameters(), lr=0.01)\n",
        "# loss_fn = torch.nn.NLLLoss(ignore_index=0)  # 0 is <PAD>\n",
        "# pbar = tqdm(range(200))\n",
        "# for i in pbar:\n",
        "#   opt.zero_grad()\n",
        "#   t = x2[:,1:].flatten(start_dim=0)\n",
        "#   y2, attn = model(x1, x2)\n",
        "#   p = y2[:,:-1,:].flatten(start_dim=0, end_dim=1)\n",
        "#   loss = loss_fn(p, t)\n",
        "#   loss.backward()\n",
        "#   opt.step()\n",
        "#   with torch.no_grad():\n",
        "#     idx = torch.nonzero(t).view(-1)\n",
        "#     acc = torch.sum(torch.argmax(p, 1)[idx] == t[idx]).float()/idx.size()[0]\n",
        "#     ps = torch.argmax(y2[:,:-1,:], 2).detach().cpu().numpy()\n",
        "#     ts = x2[:,1:].detach().cpu().numpy()\n",
        "#     for sp, st in zip(ps, ts):\n",
        "#       t_set = set(st) - {0, 1, 2, 3}\n",
        "#       eos = sp[sp == 3][0] if 3 in sp else len(sp)\n",
        "#       p_eos = sp[:eos]\n",
        "#       p_set = set(p_eos) - {0, 1, 2, 3}\n",
        "#       i_set = t_set.intersection(p_set)\n",
        "#       pre = len(i_set)/len(p_set) if len(p_set) != 0 else 0\n",
        "#       rec = len(i_set)/len(t_set) if len(t_set) != 0 else 0\n",
        "#       f1 = 2*pre*rec/(pre + rec) if pre + rec != 0 else 0\n",
        "#       # print(st, sp, p_eos)\n",
        "#       # print(t_set, p_set)\n",
        "#       # print(f'precision: {pre}, recall: {rec}, f1: {f1}')\n",
        "#     if i % 25 == 0:\n",
        "#       # print(attn)\n",
        "#       # print(p)\n",
        "#       # print(torch.argmax(p, 1))\n",
        "#       # print(t)\n",
        "#       if attn is not None:\n",
        "#          plot_attention(attn.detach().cpu().numpy(), \n",
        "#                         x1.detach().cpu().numpy(), \n",
        "#                         x2.detach().cpu().numpy(),\n",
        "#                         labels=True, tight=False,\n",
        "#                         shift=False, mask=False, figsize=(10, 5))\n",
        "#   pbar.set_description(f'loss: {loss:.3f}, acc: {acc:.3f}')\n",
        "#   if acc == 1:\n",
        "#     break\n",
        "# # Image(make_dot(attn).render('attn', format='png'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AZPHNG9iXXou",
        "colab_type": "text"
      },
      "source": [
        "# Train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bDVQqjgHhR9O",
        "colab_type": "text"
      },
      "source": [
        "## GPU/CPU stats"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PUcqiqInpBGs",
        "colab_type": "code",
        "outputId": "ae970fa3-5a45-4972-aa31-92b8840e64b1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "!pip install gputil  # for monitoring\n",
        "!pip install psutil  # for monitoring\n",
        "!pip install humanize  # for monitoring\n",
        "\n",
        "!pip install optuna  # for train\n",
        "\n",
        "!pip install rouge_score  # for metrics"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting gputil\n",
            "  Downloading https://files.pythonhosted.org/packages/ed/0e/5c61eedde9f6c87713e89d794f01e378cfd9565847d4576fa627d758c554/GPUtil-1.4.0.tar.gz\n",
            "Building wheels for collected packages: gputil\n",
            "  Building wheel for gputil (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for gputil: filename=GPUtil-1.4.0-cp36-none-any.whl size=7413 sha256=b3a5bbaaaa3f7555cfa4d3b6e6e9a8d63187571a654a3322f337932ef740721b\n",
            "  Stored in directory: /root/.cache/pip/wheels/3d/77/07/80562de4bb0786e5ea186911a2c831fdd0018bda69beab71fd\n",
            "Successfully built gputil\n",
            "Installing collected packages: gputil\n",
            "Successfully installed gputil-1.4.0\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.6/dist-packages (5.4.8)\n",
            "Requirement already satisfied: humanize in /usr/local/lib/python3.6/dist-packages (0.5.1)\n",
            "Collecting optuna\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/bc/7a/b6e3ddae75af5f7d987b4fc6be7bb38d9c9c195bd662d8762a75163c1103/optuna-1.4.0.tar.gz (183kB)\n",
            "\u001b[K     |████████████████████████████████| 184kB 6.3MB/s \n",
            "\u001b[?25hCollecting alembic\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/60/1e/cabc75a189de0fbb2841d0975243e59bde8b7822bacbb95008ac6fe9ad47/alembic-1.4.2.tar.gz (1.1MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1MB 20.8MB/s \n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting cliff\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/b9/17/57187872842bf9f65815b6969b515528ec7fd754137d2d3f49e3bc016175/cliff-3.1.0-py3-none-any.whl (80kB)\n",
            "\u001b[K     |████████████████████████████████| 81kB 6.4MB/s \n",
            "\u001b[?25hCollecting cmaes>=0.3.2\n",
            "  Downloading https://files.pythonhosted.org/packages/03/de/6ed34ebc0e5c34ed371d898540bca36edb4afe5bb2ca382483054e573c75/cmaes-0.5.0-py3-none-any.whl\n",
            "Collecting colorlog\n",
            "  Downloading https://files.pythonhosted.org/packages/00/0d/22c73c2eccb21dd3498df7d22c0b1d4a30f5a5fb3feb64e1ce06bc247747/colorlog-4.1.0-py2.py3-none-any.whl\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.6/dist-packages (from optuna) (0.15.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from optuna) (1.18.4)\n",
            "Requirement already satisfied: scipy!=1.4.0 in /usr/local/lib/python3.6/dist-packages (from optuna) (1.4.1)\n",
            "Requirement already satisfied: sqlalchemy>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from optuna) (1.3.17)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from optuna) (4.41.1)\n",
            "Collecting python-editor>=0.3\n",
            "  Downloading https://files.pythonhosted.org/packages/c6/d3/201fc3abe391bbae6606e6f1d598c15d367033332bd54352b12f35513717/python_editor-1.0.4-py3-none-any.whl\n",
            "Collecting Mako\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/a6/37/0e706200d22172eb8fa17d68a7ae22dec7631a0a92266634fb518a88a5b2/Mako-1.1.3-py2.py3-none-any.whl (75kB)\n",
            "\u001b[K     |████████████████████████████████| 81kB 8.3MB/s \n",
            "\u001b[?25hRequirement already satisfied: python-dateutil in /usr/local/lib/python3.6/dist-packages (from alembic->optuna) (2.8.1)\n",
            "Requirement already satisfied: PyYAML>=3.12 in /usr/local/lib/python3.6/dist-packages (from cliff->optuna) (3.13)\n",
            "Requirement already satisfied: pyparsing>=2.1.0 in /usr/local/lib/python3.6/dist-packages (from cliff->optuna) (2.4.7)\n",
            "Collecting cmd2!=0.8.3,<0.9.0,>=0.8.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/e9/40/a71caa2aaff10c73612a7106e2d35f693e85b8cf6e37ab0774274bca3cf9/cmd2-0.8.9-py2.py3-none-any.whl (53kB)\n",
            "\u001b[K     |████████████████████████████████| 61kB 6.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: PrettyTable<0.8,>=0.7.2 in /usr/local/lib/python3.6/dist-packages (from cliff->optuna) (0.7.2)\n",
            "Collecting pbr!=2.1.0,>=2.0.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/96/ba/aa953a11ec014b23df057ecdbc922fdb40ca8463466b1193f3367d2711a6/pbr-5.4.5-py2.py3-none-any.whl (110kB)\n",
            "\u001b[K     |████████████████████████████████| 112kB 32.6MB/s \n",
            "\u001b[?25hCollecting stevedore>=1.20.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/e6/49/a35dd566626892d577e426dbe5ea424dd7fbe10645f2c1070dcba474eca9/stevedore-1.32.0-py2.py3-none-any.whl (43kB)\n",
            "\u001b[K     |████████████████████████████████| 51kB 5.1MB/s \n",
            "\u001b[?25hRequirement already satisfied: six>=1.10.0 in /usr/local/lib/python3.6/dist-packages (from cliff->optuna) (1.12.0)\n",
            "Requirement already satisfied: MarkupSafe>=0.9.2 in /usr/local/lib/python3.6/dist-packages (from Mako->alembic->optuna) (1.1.1)\n",
            "Requirement already satisfied: wcwidth; sys_platform != \"win32\" in /usr/local/lib/python3.6/dist-packages (from cmd2!=0.8.3,<0.9.0,>=0.8.0->cliff->optuna) (0.1.9)\n",
            "Collecting pyperclip\n",
            "  Downloading https://files.pythonhosted.org/packages/f6/5b/55866e1cde0f86f5eec59dab5de8a66628cb0d53da74b8dbc15ad8dabda3/pyperclip-1.8.0.tar.gz\n",
            "Building wheels for collected packages: alembic\n",
            "  Building wheel for alembic (PEP 517) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for alembic: filename=alembic-1.4.2-cp36-none-any.whl size=159543 sha256=6b7b2b8d90605158af28dc06d6f615699713eb82bf4c891263876e09fb0d6ff6\n",
            "  Stored in directory: /root/.cache/pip/wheels/1f/04/83/76023f7a4c14688c0b5c2682a96392cfdd3ee4449eaaa287ef\n",
            "Successfully built alembic\n",
            "Building wheels for collected packages: optuna, pyperclip\n",
            "  Building wheel for optuna (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for optuna: filename=optuna-1.4.0-cp36-none-any.whl size=254554 sha256=847aee03ec936d34f0275093b18b0c8b1551f7eec28c1d9765e993b80a78a084\n",
            "  Stored in directory: /root/.cache/pip/wheels/3e/48/95/5257941c80e08f78a7afbd7981f24f9a69a83a5a1c3f160ff1\n",
            "  Building wheel for pyperclip (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyperclip: filename=pyperclip-1.8.0-cp36-none-any.whl size=8693 sha256=daee7d4b8033d38954364682fdc9824f5abc7658244e9dbea7623d13662cde6c\n",
            "  Stored in directory: /root/.cache/pip/wheels/b2/ac/0a/b784f0afe26eaf52e88a7e15c7369090deea0354fa1c6fc689\n",
            "Successfully built optuna pyperclip\n",
            "Installing collected packages: python-editor, Mako, alembic, pyperclip, cmd2, pbr, stevedore, cliff, cmaes, colorlog, optuna\n",
            "Successfully installed Mako-1.1.3 alembic-1.4.2 cliff-3.1.0 cmaes-0.5.0 cmd2-0.8.9 colorlog-4.1.0 optuna-1.4.0 pbr-5.4.5 pyperclip-1.8.0 python-editor-1.0.4 stevedore-1.32.0\n",
            "Collecting rouge_score\n",
            "  Downloading https://files.pythonhosted.org/packages/d1/6d/2b9a64cba1e4e6ecd4effbf6834b2592b54dc813654f84029758e5daeeb5/rouge_score-0.0.3-py3-none-any.whl\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.6/dist-packages (from rouge_score) (1.12.0)\n",
            "Requirement already satisfied: nltk in /usr/local/lib/python3.6/dist-packages (from rouge_score) (3.2.5)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from rouge_score) (1.18.4)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.6/dist-packages (from rouge_score) (0.9.0)\n",
            "Installing collected packages: rouge-score\n",
            "Successfully installed rouge-score-0.0.3\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "icz-Yt2dikV-",
        "colab_type": "code",
        "outputId": "2c9022c1-0b06-4357-e970-7245193ea0d3",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 318
        }
      },
      "source": [
        "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
        "import subprocess\n",
        "print(subprocess.getstatusoutput('nvidia-smi')[1])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Sat May 30 20:39:16 2020       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 418.67       Driver Version: 418.67       CUDA Version: 10.1     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   32C    P8    29W / 149W |      0MiB / 11441MiB |      0%      Default |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                       GPU Memory |\n",
            "|  GPU       PID   Type   Process name                             Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jb4ephjLpFFY",
        "colab_type": "code",
        "outputId": "74bb4edd-defe-4c3f-dee3-b2cbfa74c083",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        }
      },
      "source": [
        "# Import packages\n",
        "import os, sys, humanize, psutil, GPUtil\n",
        "\n",
        "# Define function\n",
        "def mem_report(verbose=True):\n",
        "  h = humanize.naturalsize\n",
        "  cmf = psutil.virtual_memory().available\n",
        "  cmt = psutil.virtual_memory().total\n",
        "  if verbose: \n",
        "    print(f\"CPU mem Free: {h(cmf)} / {h(cmt)}\")\n",
        "  gmfs, gmts, gmus = [], [], []\n",
        "  GPUs = GPUtil.getGPUs()\n",
        "  for i, gpu in enumerate(GPUs):\n",
        "    gmfs.append(gpu.memoryFree)\n",
        "    gmts.append(gpu.memoryTotal)\n",
        "    gmus.append(gpu.memoryUtil)\n",
        "    if verbose:\n",
        "      print(f'GPU {i} mem free: {h(gmfs[-1]*1000000)} / {h(gmts[-1]*1000000)} util: {int(gmus[-1]*100)}%')\n",
        "  return cmf, gmfs, gmts, gmus\n",
        "mem_report()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CPU mem Free: 21.5 GB / 27.4 GB\n",
            "GPU 0 mem free: 11.4 GB / 11.4 GB util: 0%\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(21480128512, [11441.0], [11441.0], [0.0])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LksNjL0ohmNR",
        "colab_type": "text"
      },
      "source": [
        "## Beam"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KcU-0oSABot9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "\n",
        "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",
        "    self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\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).to(self.device)  # 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', max_input_size=30000, \n",
        "               first_beam_width=10, verbose=0,\n",
        "               plot_dist=False, plot_size=100, **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",
        "    self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    self.max_input_size = max_input_size\n",
        "    self.first_beam_width = first_beam_width\n",
        "    self.plot_size = plot_size\n",
        "    self.plot_dist = plot_dist\n",
        "    self.verbose = verbose\n",
        "\n",
        "  def __call__(self, model, bx1, bx2=None):\n",
        "    batch_size = bx1.size(0)  # input batch_size\n",
        "    device = torch\n",
        "    bx2 = torch.full((batch_size, 1), self.sos, dtype=bx1.dtype).to(self.device)  # 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)).to(self.device)  # scores for each beam\n",
        "      bx1t = bx1 if batch_size == 1 else bx1.clone()\n",
        "      for i in range(self.beam_depth):\n",
        "        eos_mask = torch.any(bx2 == self.eos, axis=1)  # beams with <eos>\n",
        "        cur_batch_size = torch.sum(~eos_mask)  # beams without <eos>\n",
        "        cur_input_size = bx1t.numel() + bx2.numel()  # B*L1 + B*L2\n",
        "        cur_beam_width = self.beam_width if i > 0 else self.first_beam_width\n",
        "        if self.verbose:\n",
        "          print(f'{i} {bx1t.size()}, {bx2.size()}, {cur_batch_size}, {cur_input_size}')\n",
        "        if bx2.size(1) - 1 >= self.max_len or cur_batch_size == 0:\n",
        "          break\n",
        "        if cur_input_size > self.max_input_size:  # split into parts\n",
        "          n_parts = cur_input_size // self.max_input_size + 1\n",
        "          bx1t_parts = torch.chunk(bx1t[~eos_mask], n_parts)\n",
        "          bx2_parts = torch.chunk(bx2[~eos_mask], n_parts)\n",
        "          masked_by2p = torch.empty((0, cur_beam_width)).to(self.device)\n",
        "          masked_bx2 = torch.empty((0, cur_beam_width), dtype=bx2.dtype).to(self.device)\n",
        "          if self.plot_dist:\n",
        "            fig, axs = plt.subplots(ncols=2)\n",
        "          for j, (bx1t_p, bx2_p) in enumerate(zip(bx1t_parts, bx2_parts)):\n",
        "            if self.verbose:\n",
        "              print(f'{j+1}/{n_parts} {bx1t_p.size()}, {bx2_p.size()}, {bx1t_p.numel() + bx2_p.numel()}')\n",
        "            by2p_p, attn = model(bx1t_p, bx2_p)\n",
        "            masked_by2p_p, masked_bx2_p = torch.topk(by2p_p[:,-1,:], \n",
        "                                                     cur_beam_width)\n",
        "            masked_by2p = torch.cat((masked_by2p, masked_by2p_p))\n",
        "            masked_bx2 = torch.cat((masked_bx2, masked_bx2_p))\n",
        "            if self.plot_dist:\n",
        "              sns.distplot(torch.topk(by2p_p[:,-1,:], self.plot_size)[0].cpu(), \n",
        "                            hist=True, kde=False, label=f'{j+1}/{n_parts}',\n",
        "                            ax=axs[1])\n",
        "              for k in range(by2p_p.size(0)):\n",
        "                sns.distplot(torch.topk(by2p_p[k,-1,:], self.plot_size)[0].cpu(), \n",
        "                             hist=True, kde=False, ax=axs[0])\n",
        "          if self.plot_dist:\n",
        "            fig.suptitle(f'step {bx2.size(1)} top {self.plot_size} scores')\n",
        "            axs[0].set_title(f'by beam')\n",
        "            axs[1].set_title(f'by part')\n",
        "            plt.legend()\n",
        "        else:\n",
        "          by2p, attn = model(bx1t[~eos_mask], bx2[~eos_mask])  # predict\n",
        "          if i > 0 and self.beam_reduction is not None:  # beam reduction\n",
        "            if i == 1:\n",
        "              prev_prev_batch_size = batch_size\n",
        "            else:\n",
        "              prev_prev_batch_size = batch_size*self.first_beam_width*self.beam_width**(i-1)\n",
        "            by2p = by2p.view(prev_prev_batch_size, cur_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(cur_beam_width, 1, 1)  # return to prev batch size\n",
        "          if self.batch_reduction is not None:  # batch reduction\n",
        "            if i > 0:\n",
        "              prev_batch_size = batch_size*self.first_beam_width*self.beam_width**(i-1)\n",
        "            else:\n",
        "              prev_batch_size = batch_size\n",
        "            by2p = self.batch_reduction(by2p, axis=0) # cumulative batch predictions\n",
        "            masked_by2p, masked_bx2 = torch.topk(by2p[-1,:], cur_beam_width)  # beams to top k last predictions\n",
        "            masked_bx2 = masked_bx2.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "            masked_by2p = masked_by2p.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "          else:\n",
        "            if self.plot_dist:\n",
        "              fig, axs = plt.subplots(ncols=2)\n",
        "              # sns.distplot(by2p[:,-1,:], hist=True, kde=False, label=f'{bx2.size(1)}')\n",
        "              sns.distplot(torch.topk(by2p[:,-1,:], self.plot_size)[0].cpu(), \n",
        "                           hist=True, kde=False, ax=axs[1])\n",
        "              for k in range(by2p.size(0)):\n",
        "                sns.distplot(torch.topk(by2p[k,-1,:], self.plot_size)[0].cpu(), \n",
        "                             hist=True, kde=False, ax=axs[0])\n",
        "              fig.suptitle(f'step {bx2.size(1)} top {self.plot_size} scores')\n",
        "              axs[0].set_title(f'by beam')\n",
        "              axs[1].set_title(f'cumulative')\n",
        "            masked_by2p, masked_bx2 = torch.topk(by2p[:,-1,:], cur_beam_width)  # beams to top k last predictions\n",
        "        next_by2p = torch.full((bx2.size(0), cur_beam_width), \n",
        "                               float('-inf'), dtype=torch.float).to(self.device)\n",
        "        next_bx2 = torch.full((bx2.size(0), cur_beam_width), \n",
        "                              0, dtype=bx2.dtype).to(self.device)\n",
        "        next_by2p[~eos_mask] = masked_by2p  # from beams without <eos>\n",
        "        next_bx2[~eos_mask] = masked_bx2  # from beams without <eos>\n",
        "        # update bx2 and scores\n",
        "        new_batch_size = batch_size*self.first_beam_width*self.beam_width**i\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, cur_beam_width, 0)  # increase batch for new scores\n",
        "        bx2 = torch.repeat_interleave(bx2, cur_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",
        "        if batch_size == 1:\n",
        "          bx1t = bx1.expand(new_batch_size, -1)  # increase batch for new beams\n",
        "        else:\n",
        "          bx1t = torch.repeat_interleave(bx1t, cur_beam_width, 0)  # increase batch for new beams\n",
        "      if self.verbose:\n",
        "        print(f'{bx1t.size()}, {bx2.size()}, {bx1t.numel() + bx2.numel()}')\n",
        "      # print(bx2)\n",
        "      # print(beam_scores)\n",
        "      if self.depth_reduction == torch.sum:\n",
        "        pad_mask = beam_scores == float('-inf')\n",
        "        # print(pad_mask)\n",
        "        beam_scores[pad_mask] = 0\n",
        "        # print(beam_scores)\n",
        "        beam_scores = torch.sum(beam_scores, axis=1) # cumulative beams scores\n",
        "      elif self.depth_reduction == torch.mean:\n",
        "        pad_mask = beam_scores == float('-inf')\n",
        "        # print(pad_mask)\n",
        "        beam_scores[pad_mask] = 0\n",
        "        # print(beam_scores)\n",
        "        beam_scores = torch.sum(beam_scores, axis=1) # cumulative beams scores\n",
        "        # print(beam_scores)\n",
        "        beam_lens = torch.sum(~pad_mask, axis=1)\n",
        "        # print(beam_lens)\n",
        "        beam_scores = beam_scores / beam_lens\n",
        "      else:\n",
        "        self.depth_reduction(beam_scores, axis=1)\n",
        "      if self.plot_dist:\n",
        "        plt.figure()\n",
        "        sns.distplot(torch.topk(beam_scores, self.plot_size)[0].cpu(),\n",
        "                     hist=True, kde=False)\n",
        "        plt.title(f'step {bx2.size(1)} top {self.plot_size} beams scores')\n",
        "      # print(beam_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, int(bx2.size(0)/batch_size), -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",
        "    self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    if self.bkw_beam_width > self.beam_width**self.beam_depth:\n",
        "      print(f'Warning! Backward beam width {self.bkw_beam_width} greater\\\n",
        "      than beam_width**beam_depth {self.beam_width**self.beam_depth}') \n",
        "      print(f'setting backword beam width to {self.beam_width**self.beam_depth}')\n",
        "      self.bkw_beam_width = self.beam_width**self.beam_depth\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).to(self.device)  # 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)).to(self.device)  # 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": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I3dpA2Fe3BME",
        "colab_type": "text"
      },
      "source": [
        "## Metric"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tCyCP83RUbXi",
        "colab_type": "code",
        "outputId": "08795f68-51ef-4434-b68d-3203a02a495e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 692
        }
      },
      "source": [
        "from rouge_score import rouge_scorer\n",
        "from pprint import pprint\n",
        "\n",
        "def metrics(bx2, by2, eos=3):\n",
        "  scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'])\n",
        "  scores = {}\n",
        "  ps = by2.detach().cpu().numpy()\n",
        "  ts = bx2.detach().cpu().numpy()\n",
        "  for sp, st in zip(ps, ts):\n",
        "    # target to set\n",
        "    st_eos_i = np.where(st==eos)[0][0] if eos in st else len(st)\n",
        "    st_eos = st[:st_eos_i]\n",
        "    t_set = set(st_eos) - {0, 1, 2, 3}\n",
        "    # prediction to set\n",
        "    sp_eos_i = np.where(sp==eos)[0][0] if eos in sp else len(sp)\n",
        "    sp_eos = sp[:sp_eos_i]\n",
        "    p_set = set(sp_eos) - {0, 1, 2, 3}\n",
        "    # metrics\n",
        "    i_set = t_set.intersection(p_set)\n",
        "    pre = len(i_set)/len(p_set) if len(p_set) != 0 else 0\n",
        "    rec = len(i_set)/len(t_set) if len(t_set) != 0 else 0\n",
        "    f1 = 2*pre*rec/(pre + rec) if pre + rec != 0 else 0\n",
        "    rl = len(sp_eos)/len(st_eos) if len(st_eos) != 0 else 0\n",
        "    # to string\n",
        "    t_str = ' '.join(map(str, st_eos))\n",
        "    p_str = ' '.join(map(str, sp_eos))\n",
        "    rs = scorer.score(t_str, p_str)\n",
        "    print(f'st: {st}, sp: {sp}')\n",
        "    print(f'eos in st?: {eos in st}, sp_eos_i: {st_eos_i}')\n",
        "    print(f'eos in sp?: {eos in sp}, sp_eos_i: {sp_eos_i}')\n",
        "    print(f'st_eos: {st_eos}, sp_eos: {sp_eos}')\n",
        "    print(f't_set: {t_set}, p_set: {p_set}')\n",
        "    print(f't_str: \"{t_str}\", p_str: \"{p_str}\"')\n",
        "    pprint(rs)\n",
        "    print(f'precision: {pre}, recall: {rec}, f1: {f1}. rl: {rl}\\n')\n",
        "    scores.setdefault('set_precision', []).append(pre)\n",
        "    scores.setdefault('set_recall', []).append(rec)\n",
        "    scores.setdefault('set_f1', []).append(f1)\n",
        "    scores.setdefault('relen', []).append(rl)\n",
        "    mean_pre = []\n",
        "    mean_rec = []\n",
        "    mean_f1 = []\n",
        "    for k, v in rs.items():\n",
        "      scores.setdefault(f'{k}_precision', []).append(v.precision)\n",
        "      scores.setdefault(f'{k}_recall', []).append(v.recall)\n",
        "      scores.setdefault(f'{k}_f1', []).append(v.fmeasure)\n",
        "      mean_pre.append(v.precision)\n",
        "      mean_rec.append(v.recall)\n",
        "      mean_f1.append(v.fmeasure)\n",
        "    scores.setdefault('rougeM_precision', []).append(sum(mean_pre)/len(mean_pre))\n",
        "    scores.setdefault('rougeM_recall', []).append(sum(mean_rec)/len(mean_rec))\n",
        "    scores.setdefault('rougeM_f1', []).append(sum(mean_f1)/len(mean_f1))\n",
        "  return scores\n",
        "\n",
        "x2 = torch.tensor([[1, 5, 3, 4, 4, 3, 3, 3, 3], \n",
        "                   [5, 3, 0, 4, 3, 0, 0, 0, 0]])\n",
        "y2 = torch.tensor([[3, 1, 4, 3, 4, 6, 7, 4], \n",
        "                   [0, 5, 4, 3, 3, 0, 0, 3]])\n",
        "metrics(x2, y2)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "st: [1 5 3 4 4 3 3 3 3], sp: [3 1 4 3 4 6 7 4]\n",
            "eos in st?: True, sp_eos_i: 2\n",
            "eos in sp?: True, sp_eos_i: 0\n",
            "st_eos: [1 5], sp_eos: []\n",
            "t_set: {5}, p_set: set()\n",
            "t_str: \"1 5\", p_str: \"\"\n",
            "{'rouge1': Score(precision=0.0, recall=0.0, fmeasure=0.0),\n",
            " 'rouge2': Score(precision=0.0, recall=0.0, fmeasure=0.0),\n",
            " 'rougeL': Score(precision=0, recall=0, fmeasure=0)}\n",
            "precision: 0, recall: 0.0, f1: 0. rl: 0.0\n",
            "\n",
            "st: [5 3 0 4 3 0 0 0 0], sp: [0 5 4 3 3 0 0 3]\n",
            "eos in st?: True, sp_eos_i: 1\n",
            "eos in sp?: True, sp_eos_i: 3\n",
            "st_eos: [5], sp_eos: [0 5 4]\n",
            "t_set: {5}, p_set: {4, 5}\n",
            "t_str: \"5\", p_str: \"0 5 4\"\n",
            "{'rouge1': Score(precision=0.3333333333333333, recall=1.0, fmeasure=0.5),\n",
            " 'rouge2': Score(precision=0.0, recall=0.0, fmeasure=0.0),\n",
            " 'rougeL': Score(precision=0.3333333333333333, recall=1.0, fmeasure=0.5)}\n",
            "precision: 0.5, recall: 1.0, f1: 0.6666666666666666. rl: 3.0\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'relen': [0.0, 3.0],\n",
              " 'rouge1_f1': [0.0, 0.5],\n",
              " 'rouge1_precision': [0.0, 0.3333333333333333],\n",
              " 'rouge1_recall': [0.0, 1.0],\n",
              " 'rouge2_f1': [0.0, 0.0],\n",
              " 'rouge2_precision': [0.0, 0.0],\n",
              " 'rouge2_recall': [0.0, 0.0],\n",
              " 'rougeL_f1': [0, 0.5],\n",
              " 'rougeL_precision': [0, 0.3333333333333333],\n",
              " 'rougeL_recall': [0, 1.0],\n",
              " 'rougeM_f1': [0.0, 0.3333333333333333],\n",
              " 'rougeM_precision': [0.0, 0.2222222222222222],\n",
              " 'rougeM_recall': [0.0, 0.6666666666666666],\n",
              " 'set_f1': [0, 0.6666666666666666],\n",
              " 'set_precision': [0, 0.5],\n",
              " 'set_recall': [0.0, 1.0]}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zhWq-kdb3Ep7",
        "colab_type": "text"
      },
      "source": [
        "## Run"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rZe5hRHIWcZD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import humanize\n",
        "import psutil\n",
        "import GPUtil\n",
        "\n",
        "def run(train, device, model, opt, loss_fn, dl, ds, print_triples=1, \n",
        "        plot_attn=False, inference=None, metric=False, pin_memory=False,\n",
        "        max_samples=None):\n",
        "  scores = {}\n",
        "  pbar = tqdm(dl, total=max_samples)\n",
        "  batch_cnt = 0\n",
        "  if train:\n",
        "    model.train()\n",
        "    for x1, x2 in pbar:\n",
        "      opt.zero_grad()\n",
        "      x1 = x1.to(device, non_blocking=pin_memory)\n",
        "      x2 = x2.to(device, non_blocking=pin_memory)\n",
        "      t = x2[:,1:].flatten(0)\n",
        "      y2p, attn = model(x1, x2)\n",
        "      p = y2p[:,:-1,:].flatten(0, 1)\n",
        "      loss = loss_fn(p, t)\n",
        "      loss.backward()\n",
        "      opt.step()\n",
        "      scores.setdefault('loss', []).append(loss.item())\n",
        "      y2 = None\n",
        "      if metric:\n",
        "        y2 = torch.argmax(y2p[:,:-1,:], 2)\n",
        "        # idx = torch.nonzero(t).view(-1)\n",
        "        # acc = torch.sum(y2.view(-1)[idx] == t[idx]).float()/idx.size()[0]\n",
        "        # accs.append(acc.item())\n",
        "        bs = metrics(x2[:,1:], y2)  # batch scores\n",
        "        r1f = sum(bs['rouge1_f1'])/len(bs['rouge1_f1'])\n",
        "        r2f = sum(bs['rouge2_f1'])/len(bs['rouge2_f1'])\n",
        "        rlf = sum(bs['rougeL_f1'])/len(bs['rougeL_f1'])\n",
        "        sf = sum(bs['set_f1'])/len(bs['set_f1'])\n",
        "        rl = sum(bs['relen'])/len(bs['relen'])\n",
        "        for score, values in bs.items():\n",
        "          scores.setdefault(score, []).extend(values)\n",
        "        cmt = psutil.virtual_memory().total\n",
        "        cma = psutil.virtual_memory().available\n",
        "        cmu = (cmt-cma)/cmt\n",
        "        gmu = GPUtil.getGPUs()[0].memoryUtil if len(GPUtil.getGPUs()) > 0 else 0\n",
        "        pbar.set_description(f'train {int(cmu*100)}%/{int(gmu*100)}% \\\n",
        "          l: {loss.item():.3}, \\\n",
        "          sf: {int(sf*100)}, r1f: {int(r1f*100)}, r2f: {int(r2f*100)}, \\\n",
        "          rlf: {int(rlf*100)}, rl: {int(rl*100)}')\n",
        "      else:\n",
        "        pbar.set_description(f'loss: {loss.item():.3}')\n",
        "      batch_cnt += 1\n",
        "      if max_samples is not None:\n",
        "        if batch_cnt == max_samples:\n",
        "          break\n",
        "  else:\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "      for x1, x2 in pbar:\n",
        "        x1 = x1.to(device, non_blocking=pin_memory)\n",
        "        x2 = x2.to(device, non_blocking=pin_memory)\n",
        "        if inference is not None:\n",
        "          y2, attn = inference(model, x1, x2)\n",
        "        else:  # forced with loss (e.g. for training)\n",
        "          t = x2[:,1:].flatten(0)\n",
        "          y2p, attn = model(x1, x2)\n",
        "          p = y2p[:,:-1,:].flatten(0, 1)\n",
        "          loss = loss_fn(p, t)\n",
        "          y2 = torch.argmax(y2p[:,:-1,:], 2)\n",
        "          scores.setdefault('loss', []).append(loss.item())\n",
        "          pbar.set_description(f'loss: {loss.item():.3}')\n",
        "        if metric:\n",
        "          bs = metrics(x2[:,1:], y2)  # batch scores\n",
        "          for score, values in bs.items():\n",
        "            scores.setdefault(score, []).extend(values)\n",
        "          # r1f = sum(bs['rouge1_f1'])/len(bs['rouge1_f1'])\n",
        "          # r2f = sum(bs['rouge2_f1'])/len(bs['rouge2_f1'])\n",
        "          # rlf = sum(bs['rougeL_f1'])/len(bs['rougeL_f1'])\n",
        "          # sf = sum(bs['set_f1'])/len(bs['set_f1'])\n",
        "          # rl = sum(bs['relen'])/len(bs['relen'])\n",
        "          # cmt = psutil.virtual_memory().total\n",
        "          # cma = psutil.virtual_memory().available\n",
        "          # cmu = (cmt-cma)/cmt\n",
        "          # gmu = GPUtil.getGPUs()[0].memoryUtil if len(GPUtil.getGPUs()) > 0 else 0\n",
        "          # pbar.set_description(f'eval {int(cmu*100)}%/{int(gmu*100)}% \\\n",
        "          #   sf: {int(sf*100)}, r1f: {int(r1f*100)}, r2f: {int(r2f*100)}, \\\n",
        "          #   rlf: {int(rlf*100)}, rl: {int(rl*100)}')\n",
        "          # pbar.set_description(f'gmu: {int(gmu*100)}%')\n",
        "        batch_cnt += 1\n",
        "        if max_samples is not None:\n",
        "          if batch_cnt == max_samples:\n",
        "            break\n",
        "  if print_triples:\n",
        "    if y2 is None:\n",
        "      y2 = torch.argmax(y2p[:,:-1,:], 2)\n",
        "    dy2 = map(ds.decode, y2.detach().cpu().numpy())\n",
        "    dx2 = map(ds.decode, x2[:,1:].detach().cpu().numpy())\n",
        "    dx1 = map(ds.decode, x1[:,1:].detach().cpu().numpy())\n",
        "    triples = list(zip(dx1, dx2, dy2))[:print_triples]\n",
        "    for triple in triples:\n",
        "      print('')\n",
        "      print(triple[0])\n",
        "      print(triple[1])\n",
        "      print(triple[2])\n",
        "  if plot_attn:\n",
        "    plot_attention(attn.detach().cpu().numpy()[:1], \n",
        "                   x1.detach().cpu().numpy()[:1], \n",
        "                   x2.detach().cpu().numpy()[:1],\n",
        "                   decoder_x1 = ds.decode,\n",
        "                   decoder_x2 = ds.decode,\n",
        "                   figsize=(10, 10),\n",
        "                   shift=True,\n",
        "                   mask=True,\n",
        "                   suptitle=None)\n",
        "  return scores"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hziAG-KBqOO8",
        "colab_type": "code",
        "outputId": "bc7fb16c-2bd6-4e28-ce30-0cda23a568f8",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 72
        }
      },
      "source": [
        "def tensors_mem():\n",
        "  for obj in gc.get_objects():\n",
        "    try:\n",
        "      if isinstance(obj, torch.Tensor):\n",
        "      # if hasattr(obj, 'data'):\n",
        "      #   if torch.is_tensor(obj.data):\n",
        "        print(type(obj), obj.size(), obj.dtype, obj.device, obj.layout, obj.memory_format)\n",
        "    except:\n",
        "      pass\n",
        "tensors_mem()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/torch/distributed/distributed_c10d.py:102: UserWarning: torch.distributed.reduce_op is deprecated, please use torch.distributed.ReduceOp instead\n",
            "  warnings.warn(\"torch.distributed.reduce_op is deprecated, please use \"\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ohVsoShwhWws",
        "colab_type": "text"
      },
      "source": [
        "## Main"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-OElaFmqpOwK",
        "colab_type": "code",
        "outputId": "d3d34ecc-7211-4a0d-9a14-fdbd43e04d25",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000,
          "referenced_widgets": [
            "e60ac13545ee4a17a94af84bbf022c30",
            "98cab96334884deb9c3a97de6a92f213",
            "3793b053272b46789944fae86bfeeffc",
            "0e66f13265c249b6a6e2c32a75489077",
            "a42bd4c2eab348c4b9c220b928c58680",
            "45070e65cb1745489fa9e8a206d0e54a",
            "c1479450481e4a008b802dca3b353bd8",
            "fefb8162039a424690605fb25d300797"
          ]
        }
      },
      "source": [
        "import torch\n",
        "import optuna\n",
        "from tqdm.notebook import tqdm\n",
        "from pprint import pprint\n",
        "import random\n",
        "import numpy as np\n",
        "from shutil import copyfile\n",
        "import GPUtil\n",
        "from torch import optim\n",
        "from torchsummary import summary\n",
        "\n",
        "def objective(trial):\n",
        "  device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "  trial.set_user_attr('device', str(device))\n",
        "  if device.type == \"cuda\":\n",
        "    print(f'gmu: {\"% \".join(str(int(x.memoryUtil*100)) for x in GPUtil.getGPUs())}%')\n",
        "    torch.cuda.empty_cache()\n",
        "    torch.cuda.ipc_collect()\n",
        "    print(torch.cuda.memory_stats(device=device))\n",
        "    print(torch.cuda.memory_summary(device=device))\n",
        "    print(f'gmu: {\"% \".join(str(int(x.memoryUtil*100)) for x in GPUtil.getGPUs())}%')\n",
        "  opt_fn_map = {\n",
        "    'SGD': optim.SGD,\n",
        "    'Adam': optim.Adam,\n",
        "    'Adagrad': optim.Adagrad,\n",
        "    'ASGD': optim.ASGD,\n",
        "    'Adamax': optim.Adamax,\n",
        "    'SparseAdam': optim.SparseAdam,\n",
        "    'AdamW': optim.AdamW,\n",
        "    'Adadelta': optim.Adadelta,\n",
        "    'LBFGS': optim.LBFGS,\n",
        "    'RMSprop': optim.RMSprop,\n",
        "    'Rprop': optim.Rprop\n",
        "  }\n",
        "  eval_fn_map = {\n",
        "      'Beam': Beam,\n",
        "      'BeamBkw': BeamBkw,\n",
        "      'Greedy': Greedy,\n",
        "      'Forced': Forced\n",
        "  }\n",
        "  # fn = 'model.pth'\n",
        "  fn = os.path.join(root, f'model_{trial.study.study_name}_{trial.number}.pth')\n",
        "  print(fn)\n",
        "  checkpoint = None\n",
        "  checkpoint = torch.load('/content/drive/My Drive/model_bpe_4.pth',\n",
        "                          map_location=device)\n",
        "  # checkpoint = torch.load('/content/drive/My Drive/model_ria_base.pth',\n",
        "  #                         map_location=device)\n",
        "  train = False \n",
        "  # train_config = checkpoint['train_config']\n",
        "  train_config = {\n",
        "    'seed': trial.suggest_int('seed', 0, 0),\n",
        "    'n_epoches': trial.suggest_int('n_epoches', 1, 1),\n",
        "    'lr': trial.suggest_loguniform('lr', 1e-3, 1e-3),\n",
        "      # 'opt_fn': trial.suggest_categorical('opt_fn', ['SGD', 'Adam', 'Adagrad', \n",
        "      #                                           'ASGD', 'Adamax',\n",
        "      #                                           'AdamW', 'Adadelta',\n",
        "      #                                           'RMSprop', 'Rprop']),\n",
        "    'opt_fn': trial.suggest_categorical('opt_fn', ['Adam']),\n",
        "    'batch_size': trial.suggest_int('batch_size', 1, 1),\n",
        "    'weight_decay': None,\n",
        "    'pin_memory': True,\n",
        "    # 'weight_decay': trial.suggest_loguniform('weight_decay', 1e-6, 1e-6),\n",
        "    'n_x1': trial.suggest_int('n_x1', 1, 1)\n",
        "  }\n",
        "  eval_config = {\n",
        "      # 'eval_fn': trial.suggest_categorical('eval_fn', ['Beam', 'BeamBkw']),\n",
        "      # 'eval_fn': trial.suggest_categorical('eval_fn', ['Beam']),\n",
        "      'eval_fn': 'Greedy',\n",
        "      'max_samples': 1}\n",
        "  if eval_config['eval_fn'] in ['Beam', 'BeamBkw']:\n",
        "    eval_config['beam_width'] = trial.suggest_int('beam_width', 2, 2)\n",
        "    eval_config['beam_depth'] = trial.suggest_int('beam_depth', 8, 8)\n",
        "    # eval_config['beam_reduction'] = trial.suggest_categorical('beam_reduction', ['sum', 'mean', 'none']),\n",
        "    eval_config['beam_reduction'] = 'none'\n",
        "    eval_config['max_input_size'] = 35000  # BATCH*(LEN1+LEN2) ~ BATCH*(10+30)\n",
        "    eval_config['batch_reduction'] = 'none'\n",
        "    eval_config['max_len'] = 16\n",
        "    eval_config['plot_size'] = 10\n",
        "    eval_config['plot_dist'] = False\n",
        "    eval_config['verbose'] = False\n",
        "    # eval_config['first_beam_width'] = eval_config['beam_width']\n",
        "    eval_config['first_beam_width'] = eval_config['beam_width']\n",
        "    # eval_config['depth_reduction'] = trial.suggest_categorical('depth_reduction', ['sum', 'mean'])\n",
        "    eval_config['depth_reduction'] = 'mean'\n",
        "    if eval_config['eval_fn'] == 'BeamBkw':\n",
        "      eval_config['bkw_beam_width'] = trial.suggest_int('bkw_beam_width', 2, 2)\n",
        "  if checkpoint is None:\n",
        "    model_config = {\n",
        "      'enc_num_embeddings': len(ds.bpe.vocab()),\n",
        "      'enc_embedding_dim': trial.suggest_int('enc_embedding_dim', 300, 300),\n",
        "      'enc_padding_idx': 0,\n",
        "      'dec_num_embeddings': len(ds.bpe.vocab()),\n",
        "      'dec_embedding_dim': trial.suggest_int('dec_embedding_dim', 300, 300),\n",
        "      'dec_padding_idx': 0,\n",
        "      'rnn_type': trial.suggest_categorical('rnn_type', ['GRU']),\n",
        "      # 'rnn_type': trial.suggest_categorical('rnn_type', ['RNN']),\n",
        "      'hidden_size': trial.suggest_int('hidden_size', 200, 200),\n",
        "      'num_layers': trial.suggest_int('num_layers', 1, 1),\n",
        "      'bidirectional': False,\n",
        "      # 'bidirectional': trial.suggest_categorical('bidirectional', [False]),\n",
        "      # 'bidirectional': trial.suggest_categorical('bidirectional', [True]),\n",
        "      'enc_dropout': trial.suggest_uniform('enc_dropout', 0.0, 0.0),\n",
        "      'dec_dropout': trial.suggest_uniform('dec_dropout', 0.0, 0.0),\n",
        "      'enc_rnn_dropout': trial.suggest_uniform('enc_rnn_dropout', 0.0, 0.0),\n",
        "      'dec_rnn_dropout': trial.suggest_uniform('dec_rnn_dropout', 0.0, 0.0),\n",
        "      'attn_type': trial.suggest_categorical('attn_type', ['soft_dist']),\n",
        "      # 'attn_type': trial.suggest_categorical('attn_type', ['dot', 'dist', 'soft_dot', 'soft_dist', 'cos', 'soft_cos', 'none']),\n",
        "      'out_hidden': trial.suggest_int('out_hidden', 600, 600),\n",
        "      'pack': trial.suggest_categorical('pack', [True]),\n",
        "    }\n",
        "  else:\n",
        "    model_config = checkpoint['model_config'] \n",
        "  # pprint(model_config)   \n",
        "  # pprint(train_config)\n",
        "  # pprint(eval_config)\n",
        "  for k, v in model_config.items():\n",
        "    trial.set_user_attr('model_' + k, v)\n",
        "  for k, v in train_config.items():\n",
        "    trial.set_user_attr('train_' + k, v)\n",
        "  for k, v in eval_config.items():\n",
        "    trial.set_user_attr('eval_' + k, v)\n",
        "  seed = train_config['seed']\n",
        "  torch.backends.cudnn.benchmark = False\n",
        "  torch.backends.cudnn.deterministic = True\n",
        "  torch.manual_seed(seed)\n",
        "  torch.cuda.manual_seed(seed)\n",
        "  random.seed(seed)\n",
        "  np.random.seed(seed)\n",
        "  opt_fn = train_config['opt_fn']\n",
        "  lr = train_config['lr']\n",
        "  n_epoches = train_config['n_epoches']\n",
        "  batch_size = train_config['batch_size']\n",
        "  weight_decay = train_config['weight_decay']\n",
        "  pin_memory = train_config['pin_memory']\n",
        "  model = EncoderDecoder(**model_config).to(device)\n",
        "  if checkpoint is not None:\n",
        "    model.load_state_dict(checkpoint['model_state_dict'])\n",
        "  pprint(model)\n",
        "  n_params = sum(x.numel() for x in model.parameters() if x.requires_grad)\n",
        "  trial.set_user_attr('n_params', n_params)\n",
        "  if opt_fn not in ['Rprop'] and weight_decay is not None:\n",
        "    opt = opt_fn_map[opt_fn](model.parameters(), lr=lr, \n",
        "                             weight_decay=weight_decay)\n",
        "  else:\n",
        "    opt = opt_fn_map[opt_fn](model.parameters(), lr=lr)\n",
        "  loss_fn = torch.nn.NLLLoss(ignore_index=0)\n",
        "  train_len = int(0.7*len(ds))\n",
        "  test_len = int(0.2*len(ds))\n",
        "  val_len = len(ds) - train_len - test_len\n",
        "  lens = [train_len, test_len, val_len]\n",
        "  trial.set_user_attr('n_samples', len(ds))\n",
        "  trial.set_user_attr('n_samples_train', train_len)\n",
        "  trial.set_user_attr('n_samples_test', test_len)\n",
        "  trial.set_user_attr('n_samples_val', val_len)\n",
        "  train_ds, test_ds, val_ds = random_split(ds, lens)\n",
        "  train_dl = DataLoader(train_ds, batch_size=batch_size, num_workers=1, \n",
        "                        shuffle=True, drop_last=False,\n",
        "                        collate_fn=Collate(train_config['n_x1']),\n",
        "                        pin_memory=pin_memory)\n",
        "  test_dl = DataLoader(test_ds, batch_size=batch_size, num_workers=1, \n",
        "                       shuffle=False, drop_last=False,\n",
        "                       collate_fn=Collate(train_config['n_x1']),\n",
        "                       pin_memory=pin_memory)\n",
        "  val_dl = DataLoader(val_ds, batch_size=batch_size, num_workers=1, \n",
        "                      shuffle=False, drop_last=False,\n",
        "                      collate_fn=Collate(train_config['n_x1']),\n",
        "                      pin_memory=pin_memory)\n",
        "  pprint(trial.user_attrs)\n",
        "  pprint(trial.params)\n",
        "  if train:\n",
        "    pbar_epoch = tqdm(range(n_epoches))\n",
        "    for i in tqdm(pbar_epoch):\n",
        "      train_scores = run(True, device, model, opt, loss_fn, train_dl, ds, \n",
        "                          print_triples=3, \n",
        "                          metric=False,\n",
        "                          pin_memory=pin_memory,\n",
        "                          plot_attn=True if model_config['attn_type'] != 'none' else False)\n",
        "      train_loss = sum(train_scores['loss'])/len(train_scores['loss'])\n",
        "      trial.report(train_loss, step=i+1)\n",
        "      # validation\n",
        "      # eval_fn = eval_fn_map[eval_config['eval_fn']](**eval_config)\n",
        "      val_scores = run(False, device, model, opt, loss_fn, val_dl, ds,  \n",
        "                       print_triples=3, \n",
        "                       inference=None,\n",
        "                       metric=False,\n",
        "                       pin_memory=pin_memory,\n",
        "                       plot_attn=False)\n",
        "      val_loss = sum(val_scores['loss'])/len(val_scores['loss'])\n",
        "      # print(f'{type(eval_fn).__name__} p: {val_pre:.3f}, r: {val_rec:.3f}, f1: {val_f1:.3f}, rl: {val_rl:.3f}')\n",
        "      pbar_epoch.set_description(f'epoch: {i+1}, tloss: {train_loss:.3}, vloss: {val_loss:.3}')\n",
        "      # pbar_epoch.set_description(f'l: {loss:.2f}/{val_loss:.2f} \\\n",
        "      # a: {int(acc*100)}/{int(val_acc*100)} p: {int(pre*100)}/{int(val_pre*100)} \\\n",
        "      # r: {int(rec*100)}/{int(val_rec*100)} f1: {int(f1*100)}/{int(val_f1*100)}\\\n",
        "      # rl: {int(rl*100)}/{int(val_rl*100)}')\n",
        "      torch.save({'epoch': i,\n",
        "                  'model_state_dict': model.state_dict(),\n",
        "                  'optimizer_state_dict': opt.state_dict(),\n",
        "                  'tloss': train_loss,\n",
        "                  'vloss': val_loss,\n",
        "                  'model_config': model_config,\n",
        "                  'train_config': train_config,\n",
        "                  'eval_config': eval_config},\n",
        "                   fn)\n",
        "      # try:\n",
        "      #   copyfile(fn, '/content/drive/My Drive/' + fn)\n",
        "      # except Exception as e:\n",
        "      #   print(e)\n",
        "  # test\n",
        "  eval_fn = eval_fn_map[eval_config['eval_fn']](**eval_config)\n",
        "  test_scores = run(False, device, model, opt, loss_fn, test_dl, ds, \n",
        "                    print_triples=1, \n",
        "                    inference=eval_fn,\n",
        "                    metric=True,\n",
        "                    pin_memory=pin_memory,\n",
        "                    max_samples=eval_config['max_samples'],\n",
        "                    plot_attn=False)\n",
        "  scores = test_scores\n",
        "  test_scores = {}\n",
        "  for score, values in scores.items():\n",
        "    test_scores[score] = sum(values)/len(values)\n",
        "    trial.set_user_attr(score, test_scores[score])\n",
        "  pprint(test_scores)\n",
        "  return train_loss if train else test_scores['rouge1_f1']\n",
        "\n",
        "import os\n",
        "fn = 'optuna_20.db'\n",
        "if not os.path.exists(fn):\n",
        "  copyfile('/content/drive/My Drive/' + fn, fn)\n",
        "study = optuna.create_study(study_name='bpe', \n",
        "                            direction='maximize', \n",
        "                            storage=f'sqlite:///{fn}', \n",
        "                            sampler=optuna.samplers.RandomSampler(0),  # optuna.samplers.GridSampler()\n",
        "                            load_if_exists=True)\n",
        "\n",
        "def dump_optuna(study, frozen_trial):\n",
        "  fn = 'optuna_20.db'\n",
        "  print(f'saving optuna {fn}')\n",
        "  try:\n",
        "    copyfile(fn, '/content/drive/My Drive/' + fn)\n",
        "  except Exception as e:\n",
        "    print(e)\n",
        "  print('done')\n",
        "\n",
        "study.optimize(objective, n_trials=1, callbacks=[dump_optuna])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[32m[I 2020-05-30 20:44:20,484]\u001b[0m Using an existing study with name 'bpe' instead of creating a new one.\u001b[0m\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "gmu: 11%\n",
            "OrderedDict([('active.all.allocated', 1662), ('active.all.current', 0), ('active.all.freed', 1662), ('active.all.peak', 66), ('active.large_pool.allocated', 84), ('active.large_pool.current', 0), ('active.large_pool.freed', 84), ('active.large_pool.peak', 19), ('active.small_pool.allocated', 1578), ('active.small_pool.current', 0), ('active.small_pool.freed', 1578), ('active.small_pool.peak', 54), ('active_bytes.all.allocated', 3050692608), ('active_bytes.all.current', 0), ('active_bytes.all.freed', 3050692608), ('active_bytes.all.peak', 986996736), ('active_bytes.large_pool.allocated', 2978718720), ('active_bytes.large_pool.current', 0), ('active_bytes.large_pool.freed', 2978718720), ('active_bytes.large_pool.peak', 977454592), ('active_bytes.small_pool.allocated', 71973888), ('active_bytes.small_pool.current', 0), ('active_bytes.small_pool.freed', 71973888), ('active_bytes.small_pool.peak', 12328448), ('allocated_bytes.all.allocated', 3050692608), ('allocated_bytes.all.current', 0), ('allocated_bytes.all.freed', 3050692608), ('allocated_bytes.all.peak', 986996736), ('allocated_bytes.large_pool.allocated', 2978718720), ('allocated_bytes.large_pool.current', 0), ('allocated_bytes.large_pool.freed', 2978718720), ('allocated_bytes.large_pool.peak', 977454592), ('allocated_bytes.small_pool.allocated', 71973888), ('allocated_bytes.small_pool.current', 0), ('allocated_bytes.small_pool.freed', 71973888), ('allocated_bytes.small_pool.peak', 12328448), ('allocation.all.allocated', 1662), ('allocation.all.current', 0), ('allocation.all.freed', 1662), ('allocation.all.peak', 66), ('allocation.large_pool.allocated', 84), ('allocation.large_pool.current', 0), ('allocation.large_pool.freed', 84), ('allocation.large_pool.peak', 19), ('allocation.small_pool.allocated', 1578), ('allocation.small_pool.current', 0), ('allocation.small_pool.freed', 1578), ('allocation.small_pool.peak', 54), ('inactive_split.all.allocated', 676), ('inactive_split.all.current', 0), ('inactive_split.all.freed', 676), ('inactive_split.all.peak', 17), ('inactive_split.large_pool.allocated', 57), ('inactive_split.large_pool.current', 0), ('inactive_split.large_pool.freed', 57), ('inactive_split.large_pool.peak', 7), ('inactive_split.small_pool.allocated', 619), ('inactive_split.small_pool.current', 0), ('inactive_split.small_pool.freed', 619), ('inactive_split.small_pool.peak', 13), ('inactive_split_bytes.all.allocated', 502255616), ('inactive_split_bytes.all.current', 0), ('inactive_split_bytes.all.freed', 502255616), ('inactive_split_bytes.all.peak', 247997952), ('inactive_split_bytes.large_pool.allocated', 388552192), ('inactive_split_bytes.large_pool.current', 0), ('inactive_split_bytes.large_pool.freed', 388552192), ('inactive_split_bytes.large_pool.peak', 243269632), ('inactive_split_bytes.small_pool.allocated', 113703424), ('inactive_split_bytes.small_pool.current', 0), ('inactive_split_bytes.small_pool.freed', 113703424), ('inactive_split_bytes.small_pool.peak', 6997504), ('num_alloc_retries', 0), ('num_ooms', 0), ('reserved_bytes.all.allocated', 3021996032), ('reserved_bytes.all.current', 0), ('reserved_bytes.all.freed', 3021996032), ('reserved_bytes.all.peak', 1008730112), ('reserved_bytes.large_pool.allocated', 2982150144), ('reserved_bytes.large_pool.current', 0), ('reserved_bytes.large_pool.freed', 2982150144), ('reserved_bytes.large_pool.peak', 994050048), ('reserved_bytes.small_pool.allocated', 39845888), ('reserved_bytes.small_pool.current', 0), ('reserved_bytes.small_pool.freed', 39845888), ('reserved_bytes.small_pool.peak', 14680064), ('segment.all.allocated', 58), ('segment.all.current', 0), ('segment.all.freed', 58), ('segment.all.peak', 20), ('segment.large_pool.allocated', 39), ('segment.large_pool.current', 0), ('segment.large_pool.freed', 39), ('segment.large_pool.peak', 13), ('segment.small_pool.allocated', 19), ('segment.small_pool.current', 0), ('segment.small_pool.freed', 19), ('segment.small_pool.peak', 7)])\n",
            "|===========================================================================|\n",
            "|                  PyTorch CUDA memory summary, device ID 0                 |\n",
            "|---------------------------------------------------------------------------|\n",
            "|            CUDA OOMs: 0            |        cudaMalloc retries: 0         |\n",
            "|===========================================================================|\n",
            "|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |\n",
            "|---------------------------------------------------------------------------|\n",
            "| Allocated memory      |       0 B  |     941 MB |    2909 MB |    2909 MB |\n",
            "|       from large pool |       0 B  |     932 MB |    2840 MB |    2840 MB |\n",
            "|       from small pool |       0 B  |      11 MB |      68 MB |      68 MB |\n",
            "|---------------------------------------------------------------------------|\n",
            "| Active memory         |       0 B  |     941 MB |    2909 MB |    2909 MB |\n",
            "|       from large pool |       0 B  |     932 MB |    2840 MB |    2840 MB |\n",
            "|       from small pool |       0 B  |      11 MB |      68 MB |      68 MB |\n",
            "|---------------------------------------------------------------------------|\n",
            "| GPU reserved memory   |       0 B  |     962 MB |    2882 MB |    2882 MB |\n",
            "|       from large pool |       0 B  |     948 MB |    2844 MB |    2844 MB |\n",
            "|       from small pool |       0 B  |      14 MB |      38 MB |      38 MB |\n",
            "|---------------------------------------------------------------------------|\n",
            "| Non-releasable memory |       0 B  |  242185 KB |  490484 KB |  490484 KB |\n",
            "|       from large pool |       0 B  |  237568 KB |  379445 KB |  379445 KB |\n",
            "|       from small pool |       0 B  |    6833 KB |  111038 KB |  111038 KB |\n",
            "|---------------------------------------------------------------------------|\n",
            "| Allocations           |       0    |      66    |    1662    |    1662    |\n",
            "|       from large pool |       0    |      19    |      84    |      84    |\n",
            "|       from small pool |       0    |      54    |    1578    |    1578    |\n",
            "|---------------------------------------------------------------------------|\n",
            "| Active allocs         |       0    |      66    |    1662    |    1662    |\n",
            "|       from large pool |       0    |      19    |      84    |      84    |\n",
            "|       from small pool |       0    |      54    |    1578    |    1578    |\n",
            "|---------------------------------------------------------------------------|\n",
            "| GPU reserved segments |       0    |      20    |      58    |      58    |\n",
            "|       from large pool |       0    |      13    |      39    |      39    |\n",
            "|       from small pool |       0    |       7    |      19    |      19    |\n",
            "|---------------------------------------------------------------------------|\n",
            "| Non-releasable allocs |       0    |      17    |     676    |     676    |\n",
            "|       from large pool |       0    |       7    |      57    |      57    |\n",
            "|       from small pool |       0    |      13    |     619    |     619    |\n",
            "|===========================================================================|\n",
            "\n",
            "gmu: 2%\n",
            "/content/drive/My Drive/model_bpe_10.pth\n",
            "EncoderDecoder(\n",
            "  (encoder): EncoderRNN(\n",
            "    (embedding): Embedding(50000, 300, padding_idx=0)\n",
            "    (rnn): RNN(300, 200, batch_first=True)\n",
            "    (dropout): Dropout(p=0.0, inplace=False)\n",
            "  )\n",
            "  (decoder): AttentionDecoderRNN(\n",
            "    (embedding): Embedding(50000, 300, padding_idx=0)\n",
            "    (rnn): RNN(300, 200, batch_first=True)\n",
            "    (out_hidden): Linear(in_features=400, out_features=600, bias=True)\n",
            "    (out): Linear(in_features=600, out_features=50000, bias=True)\n",
            "    (dropout): Dropout(p=0.0, inplace=False)\n",
            "    (softmax): LogSoftmax()\n",
            "  )\n",
            ")\n",
            "{'device': 'cuda',\n",
            " 'eval_eval_fn': 'Greedy',\n",
            " 'eval_max_samples': 1,\n",
            " 'model_attn_type': 'soft_dist',\n",
            " 'model_bidirectional': False,\n",
            " 'model_dec_dropout': 0.0,\n",
            " 'model_dec_embedding_dim': 300,\n",
            " 'model_dec_num_embeddings': 50000,\n",
            " 'model_dec_padding_idx': 0,\n",
            " 'model_dec_rnn_dropout': 0.0,\n",
            " 'model_enc_dropout': 0.0,\n",
            " 'model_enc_embedding_dim': 300,\n",
            " 'model_enc_num_embeddings': 50000,\n",
            " 'model_enc_padding_idx': 0,\n",
            " 'model_enc_rnn_dropout': 0.0,\n",
            " 'model_hidden_size': 200,\n",
            " 'model_num_layers': 1,\n",
            " 'model_out_hidden': 600,\n",
            " 'model_pack': True,\n",
            " 'model_rnn_type': 'RNN',\n",
            " 'n_params': 60491400,\n",
            " 'n_samples': 1003723,\n",
            " 'n_samples_test': 100372,\n",
            " 'n_samples_train': 702606,\n",
            " 'n_samples_val': 200745,\n",
            " 'train_batch_size': 1,\n",
            " 'train_lr': 0.001,\n",
            " 'train_n_epoches': 1,\n",
            " 'train_n_x1': 1,\n",
            " 'train_opt_fn': 'Adam',\n",
            " 'train_pin_memory': True,\n",
            " 'train_seed': 0,\n",
            " 'train_weight_decay': None}\n",
            "{'batch_size': 1,\n",
            " 'lr': 0.001,\n",
            " 'n_epoches': 1,\n",
            " 'n_x1': 1,\n",
            " 'opt_fn': 'Adam',\n",
            " 'seed': 0}\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "e60ac13545ee4a17a94af84bbf022c30",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=1.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "st: [12488   138 20019    73 41619   767  7648  4085 28514   402 18921     3], sp: [12488   138 20019  7648 41619   767 18921     3]\n",
            "eos in st?: True, sp_eos_i: 11\n",
            "eos in sp?: True, sp_eos_i: 7\n",
            "st_eos: [12488   138 20019    73 41619   767  7648  4085 28514   402 18921], sp_eos: [12488   138 20019  7648 41619   767 18921]\n",
            "t_set: {7648, 28514, 12488, 73, 138, 18921, 402, 20019, 41619, 4085, 767}, p_set: {7648, 12488, 18921, 138, 20019, 41619, 767}\n",
            "t_str: \"12488 138 20019 73 41619 767 7648 4085 28514 402 18921\", p_str: \"12488 138 20019 7648 41619 767 18921\"\n",
            "{'rouge1': Score(precision=1.0, recall=0.6363636363636364, fmeasure=0.7777777777777778),\n",
            " 'rouge2': Score(precision=0.5, recall=0.3, fmeasure=0.37499999999999994),\n",
            " 'rougeL': Score(precision=0.8571428571428571, recall=0.5454545454545454, fmeasure=0.6666666666666665)}\n",
            "precision: 1.0, recall: 0.6363636363636364, f1: 0.7777777777777778. rl: 0.6363636363636364\n",
            "\n",
            "\n",
            "['приказ об увольнении журналиста радиостанции эхо москвы александра плющева чья запись в twitter стала причиной конфликта с представителем акционеров радиостанции будет отменен сообщил главред станции алексей венедиктов<EOS>']\n",
            "['приказ об увольнении с эха москвы журналиста плющева будет отменен<EOS>']\n",
            "['приказ об увольнении журналиста эха москвы отменен<EOS>']\n",
            "\n",
            "{'relen': 0.6363636363636364,\n",
            " 'rouge1_f1': 0.7777777777777778,\n",
            " 'rouge1_precision': 1.0,\n",
            " 'rouge1_recall': 0.6363636363636364,\n",
            " 'rouge2_f1': 0.37499999999999994,\n",
            " 'rouge2_precision': 0.5,\n",
            " 'rouge2_recall': 0.3,\n",
            " 'rougeL_f1': 0.6666666666666665,\n",
            " 'rougeL_precision': 0.8571428571428571,\n",
            " 'rougeL_recall': 0.5454545454545454,\n",
            " 'rougeM_f1': 0.6064814814814814,\n",
            " 'rougeM_precision': 0.7857142857142857,\n",
            " 'rougeM_recall': 0.4939393939393939,\n",
            " 'set_f1': 0.7777777777777778,\n",
            " 'set_precision': 1.0,\n",
            " 'set_recall': 0.6363636363636364}\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\u001b[32m[I 2020-05-30 20:44:24,693]\u001b[0m Finished trial#10 with value: 0.7777777777777778 with parameters: {'batch_size': 1, 'lr': 0.001, 'n_epoches': 1, 'n_x1': 1, 'opt_fn': 'Adam', 'seed': 0}. Best is trial#10 with value: 0.7777777777777778.\u001b[0m\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "saving optuna optuna_20.db\n",
            "done\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gIc7J6j5CZM7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import gc\n",
        "gc.collect()\n",
        "print(torch.cuda.memory_summary(device=torch.device(\"cuda\" )))\n",
        "\n",
        "# try:\n",
        "#   copyfile(fn, '/content/drive/My Drive/' + fn)\n",
        "# except Exception as e:\n",
        "#   print(e)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UlcXPvqiyHbQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Greedy 200k\n",
        "# {'relen': 0.9892607721833844,\n",
        "#  'rouge1_f1': 0.38423491185005065,\n",
        "#  'rouge1_precision': 0.4026889198348889,\n",
        "#  'rouge1_recall': 0.37827062666600875,\n",
        "#  'rouge2_f1': 0.16623936163781577,\n",
        "#  'rouge2_precision': 0.17458141088353102,\n",
        "#  'rouge2_recall': 0.16407297971294255,\n",
        "#  'rougeL_f1': 0.3567234511838257,\n",
        "#  'rougeL_precision': 0.3739773768993677,\n",
        "#  'rougeL_recall': 0.35121314215560523,\n",
        "#  'set_f1': 0.39064669071317015,\n",
        "#  'set_precision': 0.4199063773953252,\n",
        "#  'set_recall': 0.3771313731307635}\n",
        "\n",
        "# try:\n",
        "# except KeyboardInterrupt:\n",
        "#   print(e)\n",
        "# except Exception as e:\n",
        "#   print(e)\n",
        "# finally:\n",
        "#   torch.save({'epoch': i,\n",
        "#               'model_state_dict': model.state_dict(),\n",
        "#               'optimizer_state_dict': opt.state_dict(),\n",
        "#               'loss': loss,\n",
        "#               'model_config': model_config,\n",
        "#               'train_config': train_config},\n",
        "#               'model_exception.pth')\n",
        "#   try:\n",
        "#     copyfile('model.pth', '/content/drive/My Drive/' + 'model.pth')\n",
        "#   except Exception as e:\n",
        "#     print(e)\n",
        "# raise KeyboardInterrupt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BJernh1FIvU2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from pprint import pprint\n",
        "pprint(study.best_params)\n",
        "pprint(study.best_value)\n",
        "pprint(study.best_trial)\n",
        "pprint(study.direction)\n",
        "%reload_ext google.colab.data_table\n",
        "df = study.trials_dataframe()\n",
        "# df.params_eval_fn.fillna(df.user_attrs_eval_eval_fn, inplace=True)\n",
        "df"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZgmudFU2IykS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# optuna.visualization.plot_contour(study, params=['beam_width', 'beam_depth', 'first_beam_width'])\n",
        "# optuna.visualization.plot_optimization_history(study)\n",
        "optuna.visualization.plot_slice(study)\n",
        "# optuna.visualization.plot_parallel_coordinate(study)  # BUG\n",
        "# optuna.visualization.plot_intermediate_values(study)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_cgvmWMJHCV_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "import pandas as pd\n",
        "\n",
        "# max_vocab = 1000\n",
        "# with open(vocab_path) as f:\n",
        "#   vocab = json.load(f)\n",
        "# top_words = dict([(w, c) for w, c in sorted(vocab.items(), key=lambda x: x[1],\n",
        "#                                             reverse=True)][:max_vocab])\n",
        "# # sns.distplot(list(map(len, vocab.keys())), kde=False)\n",
        "# # plt.figure()\n",
        "# # sns.distplot(list(map(len, top_words.keys())), kde=True)\n",
        "# # plt.figure()\n",
        "# sns.distplot(list(vocab.values()), kde=False, hist_kws={'log': True})\n",
        "# sns.distplot(list(top_words.values()), kde=False, hist_kws={'log': True})\n",
        "# ax = sns.distplot(list(vocab.values()), kde=False, hist=True, \n",
        "#              hist_kws={'log': True}, norm_hist=False)\n",
        "# # ax.set_xscale('log')\n",
        "# ax = sns.distplot(list(top_words.values()), kde=False, hist=True, \n",
        "#              hist_kws={'log': True}, norm_hist=False)\n",
        "# # ax.set_xscale('log')\n",
        "# plt.figure()\n",
        "# plt.title('Sentences lengths')\n",
        "# sns.kdeplot([len(y) for x in ds.samples for y in x[0]], \n",
        "#             label='texts', shade=True, bw=2)\n",
        "# sns.kdeplot([len(y) for x in ds.samples for y in x[1]], \n",
        "#             label='titles', shade=True, bw=2)\n",
        "# plt.legend()\n",
        "# plt.figure()\n",
        "# sns.jointplot(x='titles lengths', y='texts lengths', data=pd.DataFrame({\n",
        "#     \"texts lengths\": [sum(map(len, x[0])) for x in ds.samples],\n",
        "#     \"titles lengths\": [sum(map(len, x[1])) for x in ds.samples]}\n",
        "# ), kind='kde', n_levels=10)\n",
        "# sns.jointplot(x='user_attrs_rouge1_recall', y='user_attrs_rouge1_precision',\n",
        "#               data=df, kind='kde', n_levels=10) \n",
        "# hue=\"sex\"\n",
        "# sns.catplot(x=\"params_beam_width\", y=\"params_beam_depth\", \n",
        "#             kind=\"violin\", hue=\"params_eval_fn\", split=True, bw=1., data=df)\n",
        "# sns.violinplot(x=df.params_beam_depth, y=df.user_attrs_rouge1_f1);\n",
        "\n",
        "# sns.catplot(x=\"params_beam_width\", y=\"user_attrs_rougeL_f1\", \n",
        "#             # hue=\"smoker\",\n",
        "#             col=\"params_eval_fn\", aspect=.6,\n",
        "#             kind=\"swarm\", data=df);\n",
        "# sns.pairplot(df, hue=\"value\", vars=[\"params_beam_width\", \"params_beam_depth\"], kind=\"kde\")\n",
        "\n",
        "# g = sns.PairGrid(df, vars=[\"params_beam_width\", \"params_beam_depth\",\n",
        "#                            \"user_attrs_rougeL_f1\", \"user_attrs_rouge2_f1\",\n",
        "#                            \"user_attrs_rouge1_f1\"])\n",
        "# g.map_upper(plt.scatter)\n",
        "# g.map_lower(sns.kdeplot)\n",
        "# g.map_diag(sns.kdeplot, lw=3, legend=False);\n",
        "\n",
        "# sns.violinplot(x=df.params_beam_depth, y=df.user_attrs_rouge1_f1);\n",
        "\n",
        "# sns.catplot(x=\"params_beam_width\", \n",
        "#             y=\"params_beam_depth\", \n",
        "#             hue=\"value\",\n",
        "#             col=\"params_eval_fn\", \n",
        "#             aspect=.618,\n",
        "#             kind=\"swarm\", palette='Blues', data=df);\n",
        "# sns.catplot(x=\"params_beam_width\", \n",
        "#             y=\"user_attrs_rouge2_f1\", \n",
        "#             hue=\"params_eval_fn\",\n",
        "#             col=\"params_beam_depth\", \n",
        "#             aspect=.618,\n",
        "#             kind=\"swarm\", data=df);\n",
        "\n",
        "# sns.lmplot(x=\"params_beam_width\", y=\"user_attrs_rouge2_f1\", \n",
        "#           #  hue=\"smoker\", \n",
        "#            col=\"params_beam_depth\", row=\"params_eval_fn\", data=df);\n",
        "\n",
        "# df.params_eval_fn.fillna(df.user_attrs_eval_params_eval_fn, inplace=True)\n",
        "\n",
        "sns.pairplot(df, x_vars=[\n",
        "                        #  \"params_beam_width\", \n",
        "                         \"params_beam_depth\",\n",
        "                        #  \"params_first_beam_width\",\n",
        "                         \"user_attrs_eval_input_limit\"\n",
        "                         ], \n",
        "             y_vars=[\"user_attrs_rouge1_f1\", \n",
        "                     \"user_attrs_rouge2_f1\",\n",
        "                     \"user_attrs_rougeL_f1\",\n",
        "                     \"user_attrs_set_f1\"],\n",
        "             hue=\"user_attrs_eval_max_samples\", \n",
        "             height=5, aspect=1., kind=\"reg\");\n",
        "\n",
        "# sns.jointplot(x='params_beam_width', y='params_beam_depth',\n",
        "#               data=df, kind='kde', n_levels=10) "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NbikIK--_CV3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}