week07_compression/train_and_export.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__this notebook__ trains a small LSTM language model and showcases its predicitons in javascript"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: CUDA_VISIBLE_DEVICES=0,1,2,3\n",
      "Name: numpy\n",
      "Version: 1.18.2\n",
      "--\n",
      "Name: tensorflow\n",
      "Version: 2.2.0\n",
      "--\n",
      "Name: subword-nmt\n",
      "Version: 0.3.7\n",
      "--\n",
      "Name: nltk\n",
      "Version: 3.5\n",
      "--\n",
      "Name: prefetch-generator\n",
      "Version: 1.0.1\n",
      "--\n",
      "Name: tensorflowjs\n",
      "Version: 2.1.0\n"
     ]
    }
   ],
   "source": [
    "%env CUDA_VISIBLE_DEVICES=0,1,2,3\n",
    "!pip show numpy tensorflow subword_nmt nltk prefetch_generator tensorflowjs | grep -A1 Name\n",
    "# note: we *need* tf2.2+, the code doesn't work on tf1.x\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers as L\n",
    "\n",
    "import nltk\n",
    "import pandas as pd\n",
    "import subword_nmt.learn_bpe, subword_nmt.apply_bpe\n",
    "from tqdm import tqdm\n",
    "from IPython.display import clear_output\n",
    "from prefetch_generator import background  # pip install prefetch_generator\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Read the data\n",
    "\n",
    "We're gonna train a model on arxiv papers based on [this dataset](https://www.kaggle.com/neelshah18/arxivdataset). We'll use the version of this dataset from [Yandex NLP course](https://github.com/yandexdataschool/nlp_course)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2020-08-05 09:15:55--  https://www.dropbox.com/s/99az9n1b57qkd9j/arxivData.json.tar.gz?dl=1\n",
      "Resolving www.dropbox.com (www.dropbox.com)... 162.125.70.1, 2620:100:6026:1::a27d:4601\n",
      "Connecting to www.dropbox.com (www.dropbox.com)|162.125.70.1|:443... connected.\n",
      "HTTP request sent, awaiting response... 301 Moved Permanently\n",
      "Location: /s/dl/99az9n1b57qkd9j/arxivData.json.tar.gz [following]\n",
      "--2020-08-05 09:15:56--  https://www.dropbox.com/s/dl/99az9n1b57qkd9j/arxivData.json.tar.gz\n",
      "Reusing existing connection to www.dropbox.com:443.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://ucc53f33f2a074d99f985994595d.dl.dropboxusercontent.com/cd/0/get/A824sc_B_ysOFhi9yvGJVbno7uQGtB6LXL3NSnWTsCM1d8nU39SYYp_GO6MOnBPKi10wCDvdla6UtXKyrPOLgfJcDcArG5B4VHwwFkJgW1ovEg/file?dl=1# [following]\n",
      "--2020-08-05 09:15:56--  https://ucc53f33f2a074d99f985994595d.dl.dropboxusercontent.com/cd/0/get/A824sc_B_ysOFhi9yvGJVbno7uQGtB6LXL3NSnWTsCM1d8nU39SYYp_GO6MOnBPKi10wCDvdla6UtXKyrPOLgfJcDcArG5B4VHwwFkJgW1ovEg/file?dl=1\n",
      "Resolving ucc53f33f2a074d99f985994595d.dl.dropboxusercontent.com (ucc53f33f2a074d99f985994595d.dl.dropboxusercontent.com)... 162.125.70.15, 2620:100:6026:15::a27d:460f\n",
      "Connecting to ucc53f33f2a074d99f985994595d.dl.dropboxusercontent.com (ucc53f33f2a074d99f985994595d.dl.dropboxusercontent.com)|162.125.70.15|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 18933283 (18M) [application/binary]\n",
      "Saving to: ‘arxivData.json.tar.gz’\n",
      "\n",
      "arxivData.json.tar. 100%[===================>]  18,06M  11,1MB/s    in 1,6s    \n",
      "\n",
      "2020-08-05 09:15:59 (11,1 MB/s) - ‘arxivData.json.tar.gz’ saved [18933283/18933283]\n",
      "\n",
      "arxivData.json\n",
      "dual recurrent attention units for visual question answering ; we propose an architecture for v@@ qa which util@@ izes recurrent layers to generate visual and textual attention . the memory character@@ istic of the proposed recurrent attention units offers a rich joint embedding of visual and textual features and enables the model to reason relations between several parts of the image and question . our single model outperforms the first place win@@ ner on the v@@ qa 1 . 0 dataset , performs within margin to the current state - of - the - art ensemble model . we also experiment with replac@@ ing attention mechanisms in other state - of - the - art models with our implementation and show increased accuracy . in both cases , our recurrent attention mechanism improves performance in tasks requiring sequential or relational reasoning on the v@@ qa dataset .\n"
     ]
    }
   ],
   "source": [
    "# Alternative manual download link: https://yadi.sk/d/_nGyU2IajjR9-w\n",
    "!wget \"https://www.dropbox.com/s/99az9n1b57qkd9j/arxivData.json.tar.gz?dl=1\" -O arxivData.json.tar.gz\n",
    "!tar -xvzf arxivData.json.tar.gz\n",
    "data = pd.read_json(\"./arxivData.json\")\n",
    "lines = data.apply(lambda row: row['title'] + ' ; ' + row['summary'], axis=1).tolist()\n",
    "\n",
    "tokenizer = nltk.tokenize.WordPunctTokenizer()\n",
    "lines = [' '.join(line).lower() for line in tokenizer.tokenize_sents(lines)]\n",
    "with open('lines.tok', 'w') as f:\n",
    "    for line in lines:\n",
    "        f.write(line + '\\n')\n",
    "        \n",
    "with open('lines.tok', 'r') as f_lines_tok, open('bpe_rules', 'w') as f_bpe:\n",
    "    subword_nmt.learn_bpe.learn_bpe(f_lines_tok, f_bpe, num_symbols=4000)\n",
    "with open('bpe_rules', 'r') as f_bpe:\n",
    "    bpeizer = subword_nmt.apply_bpe.BPE(f_bpe)\n",
    "lines = list(map(' '.join, map(bpeizer.segment_tokens, map(str.split, lines))))\n",
    "print(lines[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "90-th percentile: 333\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAX0AAAD4CAYAAAAAczaOAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAWNUlEQVR4nO3df6xc5Z3f8fdnIZCERNiAa7G2VTuKlYhUDaFX4ChRlMYbY2AV8weLiFbFjVy5ammbbCttTSsV5ZdEqtWyQWrYWsFbJ8pCWDYpFqTLuoaoaqUAlx8h/AjrG37EtgDfYEO6iTaNs9/+Mc8lE++93Ln23LnjnPdLuprnPOc5Z77jO/7Muc+cmZOqQpLUDb+x1AVIkkbH0JekDjH0JalDDH1J6hBDX5I65PSlLuCNnHfeebV27dqlLkOSTikPP/zwj6pqxWzrxjr0165dy+Tk5FKXIUmnlCQvzLXO6R1J6hBDX5I6xNCXpA4x9CWpQwx9SeoQQ1+SOsTQl6QOMfQlqUMMfUnqkLH+RK4WZu2Oe05q++dvvGJIlUgaVx7pS1KHGPqS1CGGviR1iKEvSR1i6EtShxj6ktQhA4V+kt9L8mSSJ5LcluTNSdYleSDJVJKvJzmjjT2zLU+19Wv79nN9638myaWL85AkSXOZN/STrAL+DTBRVf8AOA24BvgCcFNVvRM4Cmxrm2wDjrb+m9o4klzQtnsPsBn4UpLThvtwJElvZNDpndOBtyQ5HXgr8CLwEeDOtn43cGVrb2nLtPUbk6T1315VP6uq54Ap4OKTfwiSpEHNG/pVdQj4A+CH9ML+NeBh4NWqOtaGHQRWtfYq4EDb9lgbf25//yzbvC7J9iSTSSanp6dP5DFJkuYwyPTOcnpH6euA3wTOojc9syiqamdVTVTVxIoVs17MXZJ0ggaZ3vkt4Lmqmq6qnwPfAD4ALGvTPQCrgUOtfQhYA9DWnw280t8/yzaSpBEYJPR/CGxI8tY2N78ReAq4H7iqjdkK3NXae9oybf19VVWt/5p2ds86YD3w4HAehiRpEPN+y2ZVPZDkTuAR4BjwKLATuAe4PcnnWt+tbZNbga8mmQKO0Dtjh6p6Mskd9F4wjgHXVdUvhvx4JElvIL2D8PE0MTFRk5OTS13GKeNkv1r5ZPi1zNL4SPJwVU3Mts5P5EpShxj6ktQhhr4kdYihL0kdYuhLUocY+pLUIYa+JHWIoS9JHWLoS1KHGPqS1CGGviR1iKEvSR1i6EtShxj6ktQhhr4kdYihL0kdMsiF0d+V5LG+nx8n+VSSc5LsTbK/3S5v45Pk5iRTSR5PclHfvra28fuTbJ37XiVJi2He0K+qZ6rqwqq6EPhHwE+BbwI7gH1VtR7Y15YBLqN3/dv1wHbgFoAk5wA3AJcAFwM3zLxQSJJGY6HTOxuBH1TVC8AWYHfr3w1c2dpbgK9Uz3eAZUnOBy4F9lbVkao6CuwFNp/0I5AkDWyhoX8NcFtrr6yqF1v7JWBla68CDvRtc7D1zdX/K5JsTzKZZHJ6enqB5UmS3sjAoZ/kDOBjwJ8dv656V1cfyhXWq2pnVU1U1cSKFSuGsUtJUrOQI/3LgEeq6uW2/HKbtqHdHm79h4A1fdutbn1z9UuSRmQhof9xfjm1A7AHmDkDZytwV1//te0sng3Aa20a6F5gU5Ll7Q3cTa1PkjQipw8yKMlZwEeBf97XfSNwR5JtwAvA1a3/W8DlwBS9M30+AVBVR5J8FniojftMVR056UcgSRrYQKFfVT8Bzj2u7xV6Z/McP7aA6+bYzy5g18LLlCQNg5/IlaQOMfQlqUMMfUnqEENfkjrE0JekDjH0JalDDH1J6hBDX5I6xNCXpA4x9CWpQwb6GgaNztod9yx1CZJ+jXmkL0kdYuhLUocY+pLUIYa+JHXIQKGfZFmSO5N8P8nTSd6f5Jwke5Psb7fL29gkuTnJVJLHk1zUt5+tbfz+JFvnvkdJ0mIY9Ej/i8BfVNW7gfcCTwM7gH1VtR7Y15ahdy3d9e1nO3ALQJJzgBuAS4CLgRtmXigkSaMxb+gnORv4EHArQFX9v6p6FdgC7G7DdgNXtvYW4CvV8x1gWbtw+qXA3qo6UlVHgb3A5qE+GknSGxrkSH8dMA38SZJHk3y5XTN3ZbvgOcBLwMrWXgUc6Nv+YOubq1+SNCKDhP7pwEXALVX1PuAn/HIqB3j9urg1jIKSbE8ymWRyenp6GLuUJDWDfCL3IHCwqh5oy3fSC/2Xk5xfVS+26ZvDbf0hYE3f9qtb3yHgw8f1f/v4O6uqncBOgImJiaG8kGjxncwniZ+/8YohViLpjcx7pF9VLwEHkryrdW0EngL2ADNn4GwF7mrtPcC17SyeDcBrbRroXmBTkuXtDdxNrU+SNCKDfvfOvwa+luQM4FngE/ReMO5Isg14Abi6jf0WcDkwBfy0jaWqjiT5LPBQG/eZqjoylEchSRrIQKFfVY8BE7Os2jjL2AKum2M/u4BdCylQkjQ8fiJXkjrE0JekDjH0JalDDH1J6hBDX5I6xNCXpA4x9CWpQwx9SeoQQ1+SOsTQl6QOMfQlqUMMfUnqEENfkjrE0JekDjH0JalDDH1J6pCBQj/J80m+l+SxJJOt75wke5Psb7fLW3+S3JxkKsnjSS7q28/WNn5/kq1z3Z8kaXEs5Ej/H1fVhVU1cwWtHcC+qloP7GvLAJcB69vPduAW6L1IADcAlwAXAzfMvFBIkkbjZKZ3tgC7W3s3cGVf/1eq5zvAsiTnA5cCe6vqSFUdBfYCm0/i/iVJCzRo6Bfwl0keTrK99a2sqhdb+yVgZWuvAg70bXuw9c3V/yuSbE8ymWRyenp6wPIkSYMY6MLowAer6lCSvwfsTfL9/pVVVUlqGAVV1U5gJ8DExMRQ9ilJ6hnoSL+qDrXbw8A36c3Jv9ymbWi3h9vwQ8Cavs1Xt765+iVJIzJv6Cc5K8nbZ9rAJuAJYA8wcwbOVuCu1t4DXNvO4tkAvNamge4FNiVZ3t7A3dT6JEkjMsj0zkrgm0lmxv9pVf1FkoeAO5JsA14Arm7jvwVcDkwBPwU+AVBVR5J8FniojftMVR0Z2iORJM1r3tCvqmeB987S/wqwcZb+Aq6bY1+7gF0LL1OSNAx+IleSOsTQl6QOMfQlqUMMfUnqEENfkjrE0JekDjH0JalDDH1J6hBDX5I6xNCXpA4x9CWpQwx9SeoQQ1+SOsTQl6QOMfQlqUMGDv0kpyV5NMndbXldkgeSTCX5epIzWv+ZbXmqrV/bt4/rW/8zSS4d9oORJL2xhRzpfxJ4um/5C8BNVfVO4CiwrfVvA462/pvaOJJcAFwDvAfYDHwpyWknV74kaSEGCv0kq4ErgC+35QAfAe5sQ3YDV7b2lrZMW7+xjd8C3F5VP6uq5+hdTvHiYTwISdJgBj3S/yPg94G/bcvnAq9W1bG2fBBY1dqrgAMAbf1rbfzr/bNsI0kagXlDP8lvA4er6uER1EOS7Ukmk0xOT0+P4i4lqTMGOdL/APCxJM8Dt9Ob1vkisCzJzIXVVwOHWvsQsAagrT8beKW/f5ZtXldVO6tqoqomVqxYseAHJEma27yhX1XXV9XqqlpL743Y+6rqd4H7gavasK3AXa29py3T1t9XVdX6r2ln96wD1gMPDu2RSJLmdfr8Q+b074Hbk3wOeBS4tfXfCnw1yRRwhN4LBVX1ZJI7gKeAY8B1VfWLk7h/SdICLSj0q+rbwLdb+1lmOfumqv4G+J05tv888PmFFilJGg4/kStJHXIy0zvSUKzdcc8Jb/v8jVcMsRLp159H+pLUIYa+JHWIoS9JHWLoS1KHGPqS1CGGviR1iKEvSR1i6EtShxj6ktQhhr4kdYhfw7AITuZrBSRpMXmkL0kdYuhLUocY+pLUIYNcGP3NSR5M8t0kTyb5dOtfl+SBJFNJvp7kjNZ/ZlueauvX9u3r+tb/TJJLF+tBSZJmN8iR/s+Aj1TVe4ELgc1JNgBfAG6qqncCR4Ftbfw24Gjrv6mNI8kF9C6d+B5gM/ClJKcN88FIkt7YIBdGr6r667b4pvZTwEeAO1v/buDK1t7SlmnrNyZJ67+9qn5WVc8BU8xyuUVJ0uIZaE4/yWlJHgMOA3uBHwCvVtWxNuQgsKq1VwEHANr614Bz+/tn2ab/vrYnmUwyOT09vfBHJEma00ChX1W/qKoLgdX0js7fvVgFVdXOqpqoqokVK1Ys1t1IUict6OydqnoVuB94P7AsycyHu1YDh1r7ELAGoK0/G3ilv3+WbSRJIzDI2Tsrkixr7bcAHwWephf+V7VhW4G7WntPW6atv6+qqvVf087uWQesBx4c1gORJM1vkK9hOB/Y3c60+Q3gjqq6O8lTwO1JPgc8Ctzaxt8KfDXJFHCE3hk7VNWTSe4AngKOAddV1S+G+3AkSW9k3tCvqseB983S/yyznH1TVX8D/M4c+/o88PmFlylJGgY/kStJHWLoS1KHGPqS1CGGviR1iKEvSR1i6EtShxj6ktQhhr4kdYihL0kdYuhLUocY+pLUIYa+JHWIoS9JHWLoS1KHDPJ9+tLYWrvjnpPa/vkbrxhSJdKpwSN9SeqQQS6XuCbJ/UmeSvJkkk+2/nOS7E2yv90ub/1JcnOSqSSPJ7mob19b2/j9SbbOdZ+SpMUxyJH+MeDfVdUFwAbguiQXADuAfVW1HtjXlgEuo3f92/XAduAW6L1IADcAl9C74tYNMy8UkqTRmDf0q+rFqnqktf8vvYuirwK2ALvbsN3Ala29BfhK9XwHWJbkfOBSYG9VHamqo8BeYPNQH40k6Q0taE4/yVp618t9AFhZVS+2VS8BK1t7FXCgb7ODrW+u/uPvY3uSySST09PTCylPkjSPgUM/yduAPwc+VVU/7l9XVQXUMAqqqp1VNVFVEytWrBjGLiVJzUChn+RN9AL/a1X1jdb9cpu2od0ebv2HgDV9m69ufXP1S5JGZJCzdwLcCjxdVX/Yt2oPMHMGzlbgrr7+a9tZPBuA19o00L3ApiTL2xu4m1qfJGlEBvlw1geAfwJ8L8ljre8/ADcCdyTZBrwAXN3WfQu4HJgCfgp8AqCqjiT5LPBQG/eZqjoylEchSRrIvKFfVf8byByrN84yvoDr5tjXLmDXQgqUJA2Pn8iVpA4x9CWpQwx9SeoQQ1+SOsTQl6QOMfQlqUMMfUnqEENfkjrE0JekDjH0JalDDH1J6hBDX5I6ZJBv2ZR+ba3dcc8Jb/v8jVcMsRJpNDzSl6QOMfQlqUMGuXLWriSHkzzR13dOkr1J9rfb5a0/SW5OMpXk8SQX9W2ztY3fn2TrbPclSVpcgxzp/zdg83F9O4B9VbUe2NeWAS4D1ref7cAt0HuRAG4ALgEuBm6YeaGQJI3OvKFfVf8LOP6yhluA3a29G7iyr/8r1fMdYFm7aPqlwN6qOlJVR4G9/N0XEknSIjvROf2V7WLnAC8BK1t7FXCgb9zB1jdX/9+RZHuSySST09PTJ1ieJGk2J33KZlVVkhpGMW1/O4GdABMTE0Pb70KdzKl8kjSuTvRI/+U2bUO7Pdz6DwFr+satbn1z9UuSRuhEQ38PMHMGzlbgrr7+a9tZPBuA19o00L3ApiTL2xu4m1qfJGmE5p3eSXIb8GHgvCQH6Z2FcyNwR5JtwAvA1W34t4DLgSngp8AnAKrqSJLPAg+1cZ+pquPfHJYkLbJ5Q7+qPj7Hqo2zjC3gujn2swvYtaDqJElD5SdyJalDDH1J6hBDX5I6xNCXpA4x9CWpQ7yIinSCvACLTkUe6UtShxj6ktQhhr4kdYihL0kdYuhLUocY+pLUIYa+JHWIoS9JHeKHs6Ql4Ae7tFR+rUPf69xK0q8a+fROks1JnkkylWTHqO9fkrpspEf6SU4D/gvwUeAg8FCSPVX11CjrkE5lTg3pZIx6eudiYKqqngVIcjuwBTD0pRE42SlPXzROfaMO/VXAgb7lg8Al/QOSbAe2t8W/TvLMAPs9D/jRUCocvnGtbVzrgvGtbVzrghHVli+c0Gbj+u82rnXBydf29+daMXZv5FbVTmDnQrZJMllVE4tU0kkZ19rGtS4Y39rGtS6wthMxrnXB4tY26jdyDwFr+pZXtz5J0giMOvQfAtYnWZfkDOAaYM+Ia5Ckzhrp9E5VHUvyr4B7gdOAXVX15BB2vaDpoBEb19rGtS4Y39rGtS6wthMxrnXBItaWqlqsfUuSxozfvSNJHWLoS1KHnPKhv5Rf65BkV5LDSZ7o6zsnyd4k+9vt8tafJDe3Oh9PctEi17Ymyf1JnkryZJJPjkN9Sd6c5MEk3211fbr1r0vyQLv/r7c3+klyZlueauvXLkZdx9V4WpJHk9w9LrUleT7J95I8lmSy9Y3Lc21ZkjuTfD/J00nePw61JXlX+/ea+flxkk+NSW2/157/TyS5rf2/GM3zrKpO2R96bwb/AHgHcAbwXeCCEd7/h4CLgCf6+v4zsKO1dwBfaO3Lgf8BBNgAPLDItZ0PXNTabwf+Crhgqetr+39ba78JeKDd3x3ANa3/j4F/0dr/Evjj1r4G+PoIfq//FvhT4O62vOS1Ac8D5x3XNy7Ptd3AP2vtM4Bl41JbX42nAS/R+9DSUv8fWAU8B7yl7/n1T0f1PFv0f+xF/kW+H7i3b/l64PoR17CWXw39Z4DzW/t84JnW/q/Ax2cbN6I676L3nUdjUx/wVuARep/K/hFw+vG/V3pner2/tU9v47KINa0G9gEfAe5uAbDktTF76C/57xI4uwVYxq224+rZBPyfcaiNX34zwTnteXM3cOmonmen+vTObF/rsGqJapmxsqpebO2XgJWtvWS1tj8H30fvqHrJ62vTJ48Bh4G99P5ae7Wqjs1y36/X1da/Bpy7GHU1fwT8PvC3bfncMamtgL9M8nB6X1UCY/C7BNYB08CftCmxLyc5a0xq63cNcFtrL2ltVXUI+APgh8CL9J43DzOi59mpHvpjrXovzUt6TmyStwF/Dnyqqn7cv26p6quqX1TVhfSOqi8G3j3qGmaT5LeBw1X18FLXMosPVtVFwGXAdUk+1L9yCZ9rp9Ob4rylqt4H/ITelMk41AZAmxv/GPBnx69bitraewhb6L1g/iZwFrB5VPd/qof+OH6tw8tJzgdot4db/8hrTfImeoH/tar6xrjVV1WvAvfT+1N2WZKZDwv23/frdbX1ZwOvLFJJHwA+luR54HZ6UzxfHIfa2tEhVXUY+Ca9F8tx+F0eBA5W1QNt+U56LwLjUNuMy4BHqurltrzUtf0W8FxVTVfVz4Fv0HvujeR5dqqH/jh+rcMeYGtrb6U3lz7Tf207Q2AD8Frfn5hDlyTArcDTVfWH41JfkhVJlrX2W+i9z/A0vfC/ao66Zuq9CrivHZ0NXVVdX1Wrq2otvefSfVX1u0tdW5Kzkrx9pk1vfvoJxuC5VlUvAQeSvKt1baT3VelLXlufj/PLqZ2ZGpayth8CG5K8tf0/nfk3G83zbLHfQFnsH3rvuP8VvXnh/zji+76N3pzcz+kd8WyjN9e2D9gP/E/gnDY29C4g8wPge8DEItf2QXp/tj4OPNZ+Ll/q+oB/CDza6noC+E+t/x3Ag8AUvT/Dz2z9b27LU239O0b0u/0wvzx7Z0lra/f/3fbz5MzzfKl/l331XQhMtt/pfweWj1FtZ9E7Kj67r2/JawM+DXy//R/4KnDmqJ5nfg2DJHXIqT69I0laAENfkjrE0JekDjH0JalDDH1J6hBDX5I6xNCXpA75/9HkvJc41CSwAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "num_tokens_per_line = list(map(len, map(str.split, lines)))\n",
    "max_len = int(np.percentile(num_tokens_per_line, 90))\n",
    "plt.hist(num_tokens_per_line, bins=20);\n",
    "print(\"90-th percentile:\", max_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Vocabulary\n",
    "\n",
    "Let's define a special class that converts between text lines and tf tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "import json\n",
    "from collections import Counter\n",
    "\n",
    "class Vocab:\n",
    "    def __init__(self, tokens, bos=\"_BOS_\", eos=\"_EOS_\", unk='_UNK_'):\n",
    "        \"\"\"\n",
    "        A special class that converts lines of tokens into matrices and backwards\n",
    "        source: https://github.com/yandexdataschool/nlp_course/blob/2019/week04_seq2seq/utils.py\n",
    "        \"\"\"\n",
    "        assert all(tok in tokens for tok in (bos, eos, unk))\n",
    "        self.tokens = tokens\n",
    "        self.token_to_ix = {t:i for i, t in enumerate(tokens)}\n",
    "        self.bos, self.eos, self.unk = bos, eos, unk\n",
    "        self.bos_ix = self.token_to_ix[bos]\n",
    "        self.eos_ix = self.token_to_ix[eos]\n",
    "        self.unk_ix = self.token_to_ix[unk]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.tokens)\n",
    "\n",
    "    @classmethod\n",
    "    def from_data(cls, lines, max_tokens=None, bos=\"_BOS_\", eos=\"_EOS_\", unk='_UNK_'):\n",
    "        flat_lines = '\\n'.join(list(lines)).split()\n",
    "        tokens, counts = zip(*Counter(flat_lines).most_common(max_tokens))\n",
    "        tokens = [bos, eos, unk] + [t for t in sorted(tokens) if t not in (bos, eos, unk)]\n",
    "        return cls(tokens, bos, eos, unk)\n",
    "    \n",
    "    def save(self, path):\n",
    "        with open(path, 'w') as f:\n",
    "            json.dump((self.tokens, self.bos, self.eos, self.unk), f)\n",
    "    \n",
    "    @classmethod\n",
    "    def load(cls, path):\n",
    "        with open(path, 'r') as f:\n",
    "            return cls(*json.load(f))\n",
    "\n",
    "    def tokenize(self, string):\n",
    "        \"\"\"converts string to a list of tokens\"\"\"\n",
    "        tokens = [tok if tok in self.token_to_ix else self.unk for tok in string.split()]\n",
    "        return [self.bos] + tokens + [self.eos]\n",
    "\n",
    "    def to_matrix(self, lines, max_len=None):\n",
    "        \"\"\"\n",
    "        convert variable length token sequences into  fixed size matrix\n",
    "        example usage:\n",
    "        >>>print( as_matrix(words[:3],source_to_ix))\n",
    "        [[15 22 21 28 27 13 -1 -1 -1 -1 -1]\n",
    "         [30 21 15 15 21 14 28 27 13 -1 -1]\n",
    "         [25 37 31 34 21 20 37 21 28 19 13]]\n",
    "        \"\"\"\n",
    "        lines = list(map(self.tokenize, lines))\n",
    "        max_len = max_len or max(map(len, lines))\n",
    "        matrix = np.full((len(lines), max_len), self.eos_ix, dtype='int32')\n",
    "        for i, seq in enumerate(lines):\n",
    "            row_ix = list(map(self.token_to_ix.get, seq))[:max_len]\n",
    "            matrix[i, :len(row_ix)] = row_ix\n",
    "\n",
    "        return tf.convert_to_tensor(matrix)\n",
    "\n",
    "    def to_lines(self, matrix, crop=True):\n",
    "        \"\"\"\n",
    "        Convert matrix of token ids into strings\n",
    "        :param matrix: matrix of tokens of int32, shape=[batch,time]\n",
    "        :param crop: if True, crops BOS and EOS from line\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        lines = []\n",
    "        for line_ix in map(list,matrix):\n",
    "            if crop:\n",
    "                if line_ix[0] == self.bos_ix:\n",
    "                    line_ix = line_ix[1:]\n",
    "                if self.eos_ix in line_ix:\n",
    "                    line_ix = line_ix[:line_ix.index(self.eos_ix)]\n",
    "            line = ' '.join(self.tokens[i] for i in line_ix)\n",
    "            lines.append(line)\n",
    "        return lines\n",
    "    \n",
    "    def infer_length(self, batch_ix: tf.Tensor, dtype=tf.int32):\n",
    "        \"\"\" compute length given output indices, return int32 vector [len(batch_ix)] \"\"\"\n",
    "        is_eos = tf.cast(tf.equal(batch_ix, self.eos_ix), dtype)\n",
    "        count_eos = tf.cumsum(is_eos, axis=1, exclusive=True)\n",
    "        lengths = tf.reduce_sum(tf.cast(tf.equal(count_eos, 0), dtype), axis=1)\n",
    "        return lengths\n",
    "\n",
    "\n",
    "    def infer_mask(self, batch_ix: tf.Tensor, dtype=tf.bool):\n",
    "        \"\"\" all tokens after (but not including) first EOS are masked out \"\"\"\n",
    "        lengths = self.infer_length(batch_ix)\n",
    "        return tf.sequence_mask(lengths, maxlen=tf.shape(batch_ix)[1], dtype=dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 149), dtype=int32, numpy=\n",
       "array([[   0, 1193, 3039,  445, 3885, 1589, 4000, 2981,  314,   87, 4025,\n",
       "        2924,  285,  358, 1589, 3930, 2962, 4045, 3925, 2080, 3039, 2152,\n",
       "        3758, 1674, 4000,  299, 3675,  445,   41, 3681, 2355,  690, 2045,\n",
       "        2576, 3681, 2925, 3039,  445, 3885, 2581,  113, 3161, 2088, 1262,\n",
       "        2576, 4000,  299, 3675, 1532,  299, 1282, 3681, 2408, 3758, 3016,\n",
       "        3079,  542, 3309, 2714, 2576, 3681, 1879,  299, 2981,   41, 2661,\n",
       "        3370, 2408, 2672, 3681, 1568, 2782, 4062, 2493, 2592, 3681, 3930,\n",
       "        2962,   51,   41,   47,  977,   35, 2751, 4069, 2297, 3758, 3681,\n",
       "         962, 3482,   37, 2576,   37, 3681,   37,  379, 1319, 2408,   41,\n",
       "        4025,  264, 1446, 4067, 3099, 1954,  445, 2345, 1905, 2654, 3482,\n",
       "          37, 2576,   37, 3681,   37,  379, 2413, 4067, 2661, 1889,  299,\n",
       "        3329, 1921,  140,   41, 1905,  577,  649,   35, 2661, 3039,  445,\n",
       "        2344, 1903, 2747, 1905, 3633, 3120, 3294, 2624, 3078, 3019, 2592,\n",
       "        3681, 3930, 2962,  977,   41,    1],\n",
       "       [   0, 3294, 3325,   37, 3672,  719, 4067, 3039,  299,  907, 2502,\n",
       "        2499,   87, 3023,  338,  499, 2592,  383, 2502, 2499,   23,  286,\n",
       "        2545,   26, 1769, 3333, 2912, 3148, 1589, 3325,   37, 3672,  719,\n",
       "          41, 1811,   35, 2290, 3325, 3674, 2570,  959, 1905, 3293,   23,\n",
       "        1204,   41, 1646,   42, 3285, 1905,  113, 1173, 2624, 3922, 3658,\n",
       "         297, 1905,  113, 1094, 1646,   28,  299, 2432, 1433,  286, 2459,\n",
       "          37,  499, 3608, 1170, 2539, 2182, 3681, 2839,  659, 1118, 3325,\n",
       "        3674, 4041,  723,  113, 3541, 3291, 1324, 2597,   41, 1905, 3707,\n",
       "        4076,   35, 4025, 2856,  113, 2408,  499, 2592, 3039, 2502, 2499,\n",
       "         299,  907, 2502, 2499, 3680, 1915,  420, 3681, 2839,  659, 1118,\n",
       "        3325, 3674,   41, 2661, 2408,  149, 3482,   37, 2576,   37, 3681,\n",
       "          37,  379, 3148, 2592, 3712, 1103,  978, 1589, 1094, 1646,  159,\n",
       "        2847,   41,    1,    1,    1,    1,    1,    1,    1,    1,    1,\n",
       "           1,    1,    1,    1,    1,    1]], dtype=int32)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "voc = Vocab.from_data(lines)\n",
    "voc.to_matrix(lines[:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['dual recurrent attention units for visual question answering ; we propose an architecture for',\n",
       " 'sequential short - text classification with recurrent and convolutional neural networks ; recent approaches',\n",
       " 'multi@@ resolution recurrent neural networks : an application to dialogue response generation ; we']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "voc.to_lines(voc.to_matrix(lines[:3])[:, :15])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model & training\n",
    "Now let as train a simple LSTM language model the pre-processed data.\n",
    "\n",
    "__Note:__ we don't use validation for simplicity's sake, meaning our model probably overfits like crazy. But who cares? its a demo!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LanguageModel(keras.models.Model):\n",
    "    def __init__(self, voc, emb_size=128, hid_size=1024):\n",
    "        super().__init__()\n",
    "        self.voc = voc\n",
    "        self.emb = L.Embedding(len(voc), emb_size)\n",
    "        self.lstm = L.LSTM(hid_size, return_sequences=True, return_state=True)\n",
    "        self.logits = L.Dense(len(voc))\n",
    "        \n",
    "    def call(self, batch_ix):\n",
    "        hid_seq, last_hid, last_cell = self.lstm(self.emb(batch_ix[:, :-1]))\n",
    "        logits = self.logits(hid_seq)\n",
    "        mask = self.voc.infer_mask(batch_ix, dtype=tf.float32)\n",
    "        \n",
    "        loss_values = tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "            logits=tf.reshape(logits, [-1, logits.shape[-1]]),\n",
    "            labels=tf.reshape(batch_ix[:, 1:], [-1])\n",
    "        )\n",
    "        mean_loss = tf.reduce_sum(loss_values * tf.reshape(mask[:, 1:], tf.shape(loss_values))) \\\n",
    "                  / tf.reduce_sum(mask[:, 1:])\n",
    "        return mean_loss\n",
    "\n",
    "\n",
    "def iterate_minibatches(lines, batch_size, cycle=True, **kwargs):\n",
    "    while True:\n",
    "        lines_shuf = [lines[i] for i in np.random.permutation(len(lines))]\n",
    "        for batch_start in range(0, len(lines_shuf), batch_size):\n",
    "            yield (voc.to_matrix(lines_shuf[batch_start: batch_start + batch_size], **kwargs),) * 2\n",
    "        if not cycle:\n",
    "            break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')\n",
      "Number of devices: 4\n"
     ]
    }
   ],
   "source": [
    "with tf.distribute.MirroredStrategy().scope() as scope:\n",
    "    print('Number of devices: {}'.format(scope.num_replicas_in_sync))\n",
    "    model = LanguageModel(voc)\n",
    "    model.compile(optimizer='adam', loss=lambda _, loss: tf.reduce_mean(loss))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading pre-trained model.\n"
     ]
    }
   ],
   "source": [
    "import glob\n",
    "checkpoint_path = './checkpoints/lstm1024_emb128_bpe4000_batch256'\n",
    "\n",
    "if glob.glob(checkpoint_path + '*'):\n",
    "    print(\"Loading pre-trained model.\")\n",
    "    model.load_weights(checkpoint_path)\n",
    "else:\n",
    "    print(\"Training from scratch\")\n",
    "    model.fit(iterate_minibatches(lines, batch_size=256, max_len=max_len), \n",
    "              epochs=100, steps_per_epoch=256,\n",
    "              callbacks=[keras.callbacks.ModelCheckpoint(checkpoint_path, monitor='loss')])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make JS-compatible language model applier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 286,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jheuristic/anaconda3/envs/py38/lib/python3.8/site-packages/tensorflowjs/converters/keras_h5_conversion.py:123: H5pyDeprecationWarning: The default file mode will change to 'r' (read-only) in h5py 3.0. To suppress this warning, pass the mode you need to h5py.File(), or set the global default h5.get_config().default_file_mode, or set the environment variable H5PY_DEFAULT_READONLY=1. Available modes are: 'r', 'r+', 'w', 'w-'/'x', 'a'. See the docs for details.\n",
      "  return h5py.File(h5file)\n"
     ]
    }
   ],
   "source": [
    "# custom keras model that\n",
    "# * applies a single step of LSTM\n",
    "# * uses pure keras, no custom python code\n",
    "\n",
    "l_prev_tokens = L.Input([None], dtype='int32')\n",
    "l_prev_hid = L.Input([model.lstm.units], dtype='float32', name='previous_lstm_hid')\n",
    "l_prev_cell = L.Input([model.lstm.units], dtype='float32', name='previous_lstm_cell')\n",
    "\n",
    "l_prev_emb = model.emb(l_prev_tokens)  # [batch, emb_size]\n",
    "_, l_new_hid, l_new_cell = model.lstm(l_prev_emb, initial_state=[l_prev_hid, l_prev_cell])\n",
    "l_new_logits = model.logits(l_new_hid)\n",
    "\n",
    "model_step = keras.models.Model([l_prev_tokens, l_prev_hid, l_prev_cell],\n",
    "                                [l_new_logits, l_new_hid, l_new_cell])\n",
    "tfjs.converters.save_keras_model(model_step, './lm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 320,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<tf.Tensor: shape=(1, 4164), dtype=float32, numpy=\n",
       " array([[ 1.4643092 , -1.8286911 ,  1.2862148 , ...,  0.52897644,\n",
       "          1.7882657 ,  2.1528049 ]], dtype=float32)>,\n",
       " <tf.Tensor: shape=(1, 1024), dtype=float32, numpy=\n",
       " array([[ 5.3198510e-01,  1.5283672e-10, -7.6133746e-01, ...,\n",
       "          4.6707136e-03,  6.1378783e-01,  5.0298193e-05]], dtype=float32)>,\n",
       " <tf.Tensor: shape=(1, 1024), dtype=float32, numpy=\n",
       " array([[ 0.7669934 ,  0.00931896, -0.9993891 , ...,  0.04036063,\n",
       "          0.7149762 ,  1.000839  ]], dtype=float32)>]"
      ]
     },
     "execution_count": 320,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test model step from python\n",
    "h = c = tf.ones([1, model.lstm.units])\n",
    "model_step((tf.convert_to_tensor([[3]], dtype='int32'), h, c))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 328,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save bpe and vocabulary\n",
    "with open('./frontend/voc.json', 'w') as f:\n",
    "    packed_bpe_rules = list(map(list, sorted(bpeizer.bpe_codes.keys(), key=bpeizer.bpe_codes.get)))\n",
    "    json.dump([model.lstm.units, model.emb.output_dim, model.logits.units,\n",
    "               packed_bpe_rules, voc.tokens, voc.bos, voc.eos, voc.unk], f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 314,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([1, 3])"
      ]
     },
     "execution_count": 314,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "voc.to_matrix(['deep neural'])[:, :-1].shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py38",
   "language": "python",
   "name": "py38"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}