tensor_poet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/domschl/tensor-poet/blob/master/tensor_poet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fEXNOWhCEAPk"
   },
   "source": [
    "# Tensor-Poet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DabS0VZ-1Zp0"
   },
   "source": [
    "Please review [ml-indie-tools](https://github.com/domschl/ml-indie-tools), a collection machine learning tools that provides support for more environment indepent code. It will access your Google Drive when using with Google Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jtpy59Yq-Qfz"
   },
   "outputs": [],
   "source": [
    "!pip install -U ml-indie-tools==0.1.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EgLLjG4yQtft"
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import sys\n",
    "import copy\n",
    "import json\n",
    "import time\n",
    "import datetime\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "from tensorflow.keras import layers, regularizers\n",
    "\n",
    "import tensorflow_datasets as tfds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U5T4m6earb1e"
   },
   "outputs": [],
   "source": [
    "from ml_indie_tools.env_tools import MLEnv\n",
    "from ml_indie_tools.Gutenberg_Dataset import Gutenberg_Dataset\n",
    "from ml_indie_tools.Text_Dataset import Text_Dataset\n",
    "\n",
    "from ml_indie_tools.keras_custom_layers import MultiHeadSelfAttention, PositionalEncoding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VmWbteSFQtfq"
   },
   "source": [
    "## Preliminary\n",
    "\n",
    "A tensorflow deep LSTM model for text generation\n",
    "\n",
    "This code can use either CPU, GPU, TPU when running on Google Colab.\n",
    "\n",
    "Select the corresponding runtime (menu: **`Runtime / Change runtime type`**)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sfZg31sMEAP1"
   },
   "source": [
    "## 0. Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "llPw84PkEAP2"
   },
   "outputs": [],
   "source": [
    "ml_env = MLEnv(platform='tf', accelerator='fastest', old_disable_eager=True)  # TODO: move to tf.function()\n",
    "ml_env.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oZ6t9b6ZwSxi"
   },
   "outputs": [],
   "source": [
    "if ml_env.is_tpu is True:\n",
    "    tpu_strategy = ml_env.tpu_strategy\n",
    "    tpu_is_init=True\n",
    "    use_eager=False\n",
    "else:\n",
    "    use_eager=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "t-TP3Pnsrb1f"
   },
   "outputs": [],
   "source": [
    "project_name='women_writers'\n",
    "use_selfattention = False  # To explore self-attention, also check out the newer https://github.com/domschl/transformer-poet\n",
    "if use_selfattention is True:\n",
    "    model_name='mhsa_v1_tf'\n",
    "else:\n",
    "    model_name='lstm_v1_tf'\n",
    "\n",
    "# NOTICE: This will request access to Google Drive, if running on Google Colab. Google Drive is used to store snapshots\n",
    "# training data. See project ml-indie-tools: https://github.com/domschl/ml-indie-tools \n",
    "root_path, project_path, model_path, data_path, log_path = ml_env.init_paths(project_name=project_name, model_name=model_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AIkcYcEuQtfx"
   },
   "source": [
    "##  1. Text library\n",
    "\n",
    "`Text_Dataset` and `Gutenberg_Dataset` classes: libraries for training, \n",
    "encoding, batch generation, and formatted source display. It read some \n",
    "books from Project Gutenberg and supports creation of training batches. \n",
    "The output functions support highlighting to allow to compare generated \n",
    "texts with the actual sources to help to identify identical (memorized) \n",
    "parts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HjkelBcNO5WV"
   },
   "outputs": [],
   "source": [
    "use_dark_mode=True # Set to false for white background. HTML-text-compare uses background-colorization to identify different sources. Those background colors are dependent on the theme type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BF8eyWnCrb1h"
   },
   "outputs": [],
   "source": [
    "logging.basicConfig(level=logging.INFO)\n",
    "cache_dir = os.path.join(data_path, 'gutenberg_cache')\n",
    "gd = Gutenberg_Dataset(root_url='https://www.gutenberg.org/dirs/', cache_dir=cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C66X7ynnrb1h"
   },
   "outputs": [],
   "source": [
    "# sample searches\n",
    "search_spec= {\"author\": [\"brontë\",\"Jane Austen\", \"Virginia Woolf\"], \"language\": [\"english\"]}\n",
    "\n",
    "book_list=gd.search(search_spec)\n",
    "book_cnt = len(book_list)\n",
    "print(f\"{book_cnt} matching books found with search {search_spec}.\")\n",
    "if book_cnt<40:\n",
    "    # Note: please verify that book_cnt is 'reasonable'. If you plan to use a large number of texts, \n",
    "    # consider [mirroring Gutenberg](https://github.com/domschl/ml-indie-tools#working-with-a-local-mirror-of-project-gutenberg)\n",
    "    book_list = gd.insert_book_texts(book_list, download_count_limit=book_cnt)  \n",
    "else:\n",
    "    logging.error(\"Please verify your book_list, a large number of books is scheduled for download. ABORTED.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MH6_7IU3upOd"
   },
   "outputs": [],
   "source": [
    "for i in range(len(book_list)):\n",
    "    print(f\"{i}: {book_list[i]['title']} - {book_list[i]['author']}, {book_list[i]['ebook_id']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2jBH3Z15rb1h"
   },
   "outputs": [],
   "source": [
    "select = (17,12,21) # 11,20,21)\n",
    "sub_book_list = [book_list[i] for i in range(len(book_list)) if i in select]\n",
    "\n",
    "print(\"Using books:\")\n",
    "for i in range(len(sub_book_list)):\n",
    "    print(f\"{i+1}: {sub_book_list[i]['title']} - {sub_book_list[i]['author']}\")\n",
    "\n",
    "td = Text_Dataset(sub_book_list)\n",
    "td.init_tokenizer(tokenizer='char')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f7_tc2Lirb1i"
   },
   "outputs": [],
   "source": [
    "SEQUENCE_LEN = 80\n",
    "\n",
    "td.init_getitem(sample_type='chargen_single_encoded', sample_length=SEQUENCE_LEN, content_stepping=1)\n",
    "num_records = len(td)\n",
    "\n",
    "print(f\"{num_records} records\")\n",
    "\n",
    "def get_sample_batch(td, batch_size, length, random_index=True):\n",
    "    for i in range(batch_size):\n",
    "        if random_index is True:\n",
    "            ind = random.randint(0, num_records-1)\n",
    "        else:\n",
    "            ind = i * td.getitem_content_stepping\n",
    "        Xi = td[ind]\n",
    "        yi = [Xi[-1]]\n",
    "        Xi[-1]=td.c2i['␚']  # use 'SUB'-stitut glyph to mark last char of input\n",
    "        if i==0:\n",
    "            smpX=np.array(Xi, dtype=np.float32)\n",
    "            smpy=np.array(yi, dtype=np.int32)\n",
    "        else:\n",
    "            smpX = np.vstack((smpX, np.array(Xi, dtype=np.float32)))\n",
    "            smpy = np.vstack((smpy, np.array(yi, dtype=np.int32)))\n",
    "    return np.array(smpX), np.array(smpy)\n",
    "\n",
    "def get_random_onehot_sample_batch(td, batch_size, length):\n",
    "    X, y = get_random_sample_batch(td, batch_size, length)\n",
    "    xoh = tf.keras.backend.one_hot(X, len(td.i2c))\n",
    "    yk = tf.keras.backend.constant(y)\n",
    "    return xoh, yk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TI3Fx6bNuR9A"
   },
   "outputs": [],
   "source": [
    "test_x, test_y = get_sample_batch(td, 5, 80, random_index=True)\n",
    "for i in range(len(test_x)):\n",
    "    print(f\"[{i}]: X=>{td.decode(test_x[i])}<, y=>{td.decode(test_y[i])}<\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "30hi0UPtEAQG"
   },
   "source": [
    "## 2. Use tf.data for texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "znpIUA3ig3gO"
   },
   "outputs": [],
   "source": [
    "if use_selfattention is False:\n",
    "    params = { # LSTM\n",
    "        'sequence_len': SEQUENCE_LEN,\n",
    "        'sample_every_n_epochs': 5,\n",
    "        'use_attention': False,\n",
    "\n",
    "        'lstm_layers': 3,\n",
    "        'lstm_units': 256,\n",
    "\n",
    "        'batch_size': 256,\n",
    "        'vocab_size': len(td.i2c),\n",
    "        'embedding_dim': len(td.i2c),\n",
    "        'learning_rate': 0.0002,\n",
    "        'clipvalue': None,\n",
    "    }\n",
    "else:\n",
    "    params = { # Multi-head self-attention\n",
    "        'sequence_len': SEQUENCE_LEN,\n",
    "        'sample_every_n_epochs': 5,\n",
    "        'use_attention': True,\n",
    "\n",
    "        'mhsa_layers': 4,\n",
    "        'heads': 2,\n",
    "        'units': 256, # len(td.i2c),\n",
    "        'norm': 'softmax',\n",
    "        'mh_normalize': True,\n",
    "        'final_relu': False,\n",
    "        'l2_regularizer': 1e-9,\n",
    "        'sa_l2_regularizer': 1e-10,\n",
    "        'sa_dropout': 0.8,       # no dropout: 0.0\n",
    "\n",
    "        'batch_size': 128,\n",
    "        'vocab_size': len(td.i2c),\n",
    "        'embedding_dim': len(td.i2c),\n",
    "        'learning_rate': 0.0005,\n",
    "        'clipvalue': None,\n",
    "    }\n",
    "\n",
    "if ml_env.is_tpu:\n",
    "    if use_selfattention is False:\n",
    "        params['learning_rate'] = 0.001  # LSTMs on current TPUv2 are *very* fragile.\n",
    "    # params['clipvalue'] = 1.0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jY3hUuhQYzdT"
   },
   "outputs": [],
   "source": [
    "num_batches = num_records // params['batch_size']\n",
    "print(f\"num_batches = {num_batches}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EeB7jugCV4lI"
   },
   "outputs": [],
   "source": [
    "def make_tf_dataset(num, random_index=False):\n",
    "    dx=[]\n",
    "    dy=[]\n",
    "    num_batches_active = num\n",
    "    for i in range(num_batches_active):\n",
    "        x,y=get_sample_batch(td, params['batch_size'], params['sequence_len'], random_index=random_index)\n",
    "        if i<1:\n",
    "            print(f\"[{num} x]: {x.shape} -> {y.shape}\")\n",
    "        dx.append(x)\n",
    "        dy.append(y)\n",
    "    # if ml_env.is_tpu is False:\n",
    "    dx=np.array(dx)\n",
    "    dy=np.array(dy)\n",
    "    data_xy = (dx, dy)\n",
    "    # print(f\"Shape of data_xy: {np.array(data_xy).shape}\")\n",
    "    tf_dataset=tf.data.Dataset.from_tensor_slices(data_xy)\n",
    "    return tf_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DCy7WmQyS9T-"
   },
   "outputs": [],
   "source": [
    "MAX_NUM_BATCHES = 8000\n",
    "\n",
    "if num_batches>MAX_NUM_BATCHES:\n",
    "    restricted_batches=MAX_NUM_BATCHES\n",
    "else:\n",
    "    restricted_batches=num_batches\n",
    "textlib_dataset = make_tf_dataset(restricted_batches, random_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "boow8wR7sLwi"
   },
   "outputs": [],
   "source": [
    "shuffle_buffer=10000\n",
    "if ml_env.is_tpu is True:\n",
    "    dataset=textlib_dataset.shuffle(shuffle_buffer).repeat()  # Otherwise TPU may run dry\n",
    "else:\n",
    "    dataset=textlib_dataset.shuffle(shuffle_buffer)  \n",
    "dataset.take(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "B-G5HLMqqbeT"
   },
   "outputs": [],
   "source": [
    "validation_dataset = make_tf_dataset(10, random_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h2n7-VcjTe8r"
   },
   "outputs": [],
   "source": [
    "def model_lstm(inputs, params):\n",
    "    lstms=[]\n",
    "    for i in range(params['lstm_layers']):\n",
    "        if i==params['lstm_layers']-1:\n",
    "            ret_seq=False\n",
    "        else:\n",
    "            ret_seq=True \n",
    "        if i==0:\n",
    "            lstms.append(layers.LSTM(params['lstm_units'], return_sequences=ret_seq, \n",
    "                                    batch_input_shape=[params['batch_size'], params['sequence_len'], params['embedding_dim']]))\n",
    "        else:\n",
    "            lstms.append(layers.LSTM(params['lstm_units'], return_sequences=ret_seq))\n",
    "    dense = layers.Dense(params['vocab_size'], activation=None)   # softmax in loss!\n",
    "    fl = layers.Flatten()\n",
    "    x = tf.one_hot(tf.cast(inputs,dtype=tf.int32), params['vocab_size'], axis=-1)\n",
    "    for i in range(params['lstm_layers']):\n",
    "        x = lstms[i](x)\n",
    "    x = dense(fl(x))\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZAzFlCVBiL0Q"
   },
   "outputs": [],
   "source": [
    "def model_mhsa(inputs, params):\n",
    "    dense = layers.Dense(params['vocab_size'], kernel_regularizer=regularizers.l2(params['l2_regularizer']))  # using softmax here prevents temperature adjust, affects 'from_logits' param in sparse_categorical loss \n",
    "    fl = layers.Flatten()\n",
    "    dr = layers.Dropout(params['sa_dropout'])\n",
    "    pe = PositionalEncoding(amplitude=0.3)\n",
    "    mhsa=[]\n",
    "    for i in range(params['mhsa_layers']):\n",
    "        mhsa.append(MultiHeadSelfAttention(params['heads'], units=params['units'], norm=params['norm'], mh_normalize=params['mh_normalize']))\n",
    "        # mhsa.append(keras.layers.MultiHeadAttention(num_heads=params['heads'], key_dim=params['units'], kernel_regularizer=regularizers.l2(params['sa_l2_regularizer'])))\n",
    "    # dense = layers.Dense(params['vocab_size'], activation=None)   # softmax in loss!\n",
    "    x = tf.one_hot(tf.cast(inputs,dtype=tf.int32), params['vocab_size'], axis=-1)\n",
    "    x = pe(x)\n",
    "    for i in range(params['mhsa_layers']):\n",
    "        x = mhsa[i](x)\n",
    "        # x = mhsa[i](x,x)\n",
    "    if params['sa_dropout']>0.0:\n",
    "        x = dr(x)\n",
    "    x = dense(fl(x))\n",
    "    return x "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4J13Gp_hjqqn"
   },
   "outputs": [],
   "source": [
    "def generate(text, model, gen_len=64, temperature=0.9, has_softmax=False, silent=False):\n",
    "    full=text[:-1]\n",
    "    lf=0\n",
    "    gen_text=\"\"\n",
    "    if silent is False:\n",
    "        print(f\"[{full}]\", end='')\n",
    "    tex=copy.copy(text)\n",
    "    if len(tex)>params['sequence_len']:\n",
    "        tex=tex[-params['sequence_len']:]\n",
    "    while len(tex) < params['sequence_len']:\n",
    "        tex=' '+tex\n",
    "    tex=tex[1:]+'␚'\n",
    "    for i in range(gen_len):\n",
    "        # print(f\"[{i}]: {tex}\")\n",
    "        input = np.array([td.encode(tex)])\n",
    "        pred = model.predict(input, batch_size=1)\n",
    "        if has_softmax is False:\n",
    "            pred /= temperature\n",
    "            # pred = tf.keras.layers.Softmax()(pred).numpy()\n",
    "            pred = tf.keras.layers.Softmax()(pred)\n",
    "            if use_eager is True:\n",
    "                pred=pred.numpy()\n",
    "            else:\n",
    "                pred=tf.keras.backend.eval(pred)\n",
    "        ci=np.random.choice(list(range(len(pred[0]))), p=pred[0]) # np.argmax(pred[0])\n",
    "        c=td.i2c[ci]\n",
    "        gen_text += c\n",
    "        if c=='\\n':\n",
    "            lf=0\n",
    "        else:\n",
    "            lf += 1\n",
    "            if (lf>80 and c==' ') or lf>120:\n",
    "                lf=0\n",
    "                gen_text+='\\n'\n",
    "        full+=c\n",
    "        tex=tex[:-1]+c+'␚'\n",
    "        tex=tex[-params['sequence_len']:]\n",
    "    if silent is False:\n",
    "        td.source_highlight(gen_text, min_quote_size=8)\n",
    "    return pred, gen_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fmFaOY_iJlP1"
   },
   "outputs": [],
   "source": [
    "if ml_env.is_tpu:\n",
    "    # Otherwise it explodes:\n",
    "    tf.compat.v1.experimental.output_all_intermediates(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nf-NHZ326NqJ"
   },
   "outputs": [],
   "source": [
    "inputs = keras.Input(shape=(params['sequence_len'],))\n",
    "if params['use_attention'] is True:\n",
    "    outputs = model_mhsa(inputs, params)\n",
    "    model = keras.Model(inputs=inputs, outputs=outputs, name=\"mhsa_v1_tf\")\n",
    "else:\n",
    "    outputs = model_lstm(inputs, params)\n",
    "    model = keras.Model(inputs=inputs, outputs=outputs, name=\"lstm_v1_tf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0t5JWEdYZNGz"
   },
   "outputs": [],
   "source": [
    "kscc = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "\n",
    "def loss(labels, logits):\n",
    "  vl=kscc(labels, logits)\n",
    "  return vl\n",
    "\n",
    "# def scalar_loss(labels, logits):\n",
    "#     vl = loss(labels ,logits)\n",
    "#     l = tf.reduce_mean(vl, axis=-1)\n",
    "#     return l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jc2kbGoAZXHi"
   },
   "outputs": [],
   "source": [
    "if params['clipvalue'] is not None:\n",
    "    opti = tf.keras.optimizers.Adam(params['learning_rate'], params['clipvalue'])\n",
    "else:\n",
    "    opti = tf.keras.optimizers.Adam(params['learning_rate'])\n",
    "\n",
    "if ml_env.is_tpu is True:\n",
    "    model.compile(optimizer=opti, loss=loss, metrics=[])\n",
    "else:\n",
    "    model.compile(optimizer=opti, loss=loss, metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8vxZF0wOEAQr"
   },
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "soB-Q8YXvndE"
   },
   "source": [
    "### Loss function, optimizer, tensorboard output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OHZurM5ei95K"
   },
   "outputs": [],
   "source": [
    "class GeneratorCallback(keras.callbacks.Callback):\n",
    "#    def on_test_end(self, logs=None):\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        # if use_selfattention is True and epoch % params['sample_every_n_epochs'] == 0:\n",
    "        if epoch % params['sample_every_n_epochs'] == 0:\n",
    "            idx=random.randint(0,len(td)-1)\n",
    "            text=td.decode(td[idx])\n",
    "            print()\n",
    "            if ml_env.is_tpu is True:\n",
    "                temp_list=[0.7]\n",
    "                gen_len=64\n",
    "            else:\n",
    "                temp_list=[0.5, 0.7, 0.9]\n",
    "                gen_len=192\n",
    "            for temp in temp_list:\n",
    "                print(f\"---------------- T={temp} ---------------\")\n",
    "                if ml_env.is_tpu is True:\n",
    "                    with tf.device('/cpu:0'):\n",
    "                        generate(text, model=model, gen_len=gen_len, temperature=temp)\n",
    "                else:\n",
    "                    generate(text, model=model, gen_len=gen_len, temperature=temp)\n",
    "            print(\"--------------------------------------\")\n",
    "\n",
    "generator_callback=GeneratorCallback()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5SKvObcsEAQ5"
   },
   "outputs": [],
   "source": [
    "# Directory where the checkpoints will be saved\n",
    "checkpoint_dir = os.path.join(model_path, 'training_checkpoints')\n",
    "# Name of the checkpoint files\n",
    "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt_{epoch}\")\n",
    "\n",
    "checkpoint_callback=tf.keras.callbacks.ModelCheckpoint(\n",
    "    filepath=checkpoint_prefix,\n",
    "    save_weights_only=True)\n",
    "\n",
    "logdir = os.path.join(log_path, datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))\n",
    "tensorboard_callback = tf.keras.callbacks.TensorBoard(logdir, update_freq='epoch', histogram_freq=0, profile_batch=0) # update_freq='epoch', "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o0Ew6pgWzeFj"
   },
   "outputs": [],
   "source": [
    "%tensorboard --logdir logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kDFbZcN0vxOB"
   },
   "source": [
    "## The actual training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kh2yUKBoEAQ8"
   },
   "outputs": [],
   "source": [
    "EPOCHS=1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sqcjnkKRuU3M"
   },
   "outputs": [],
   "source": [
    "restricted_batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RLbsTmtnEAQ-"
   },
   "outputs": [],
   "source": [
    "if ml_env.is_tpu is True:\n",
    "    steps_per_epoch=restricted_batches//params['batch_size']\n",
    "    if steps_per_epoch==0:\n",
    "        steps_per_epoch=1\n",
    "    history = model.fit(dataset, epochs=EPOCHS, steps_per_epoch=steps_per_epoch, callbacks=[checkpoint_callback, tensorboard_callback, generator_callback])\n",
    "else:\n",
    "    history = model.fit(dataset, validation_data=validation_dataset, epochs=EPOCHS, callbacks=[checkpoint_callback, tensorboard_callback, generator_callback])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8Z0uU33ANHPP"
   },
   "source": [
    "\n",
    "## Generate text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kvTAx5pesz8i"
   },
   "outputs": [],
   "source": [
    "model_params_gen = copy.copy(params)\n",
    "model_params_gen['batch_size'] = 1\n",
    "model_params_gen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k__-ihyvwClt"
   },
   "outputs": [],
   "source": [
    "inputs_gen = keras.Input(shape=(SEQUENCE_LEN,))\n",
    "if params['use_attention'] is True:\n",
    "    outputs_gen = model_mhsa(inputs_gen, model_params_gen)\n",
    "else:\n",
    "    outputs_gen = model_lstm(inputs_gen, model_params_gen)\n",
    "print(f\"{inputs_gen.shape} -> {outputs_gen.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5uKKo0RQs_OX"
   },
   "outputs": [],
   "source": [
    "if params['use_attention'] is True:\n",
    "    model_gen = keras.Model(inputs=inputs_gen, outputs=outputs_gen, name=\"mhsa_v1_tf_gen\")\n",
    "else:\n",
    "    model_gen = keras.Model(inputs=inputs_gen, outputs=outputs_gen, name=\"lstm_v1_tf_gen\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KCunzXAQKBI6"
   },
   "outputs": [],
   "source": [
    "checkpoint_dir = os.path.join(model_path, 'training_checkpoints')\n",
    "last_checkpoint = tf.train.latest_checkpoint(checkpoint_dir)\n",
    "print(f\"Last checkpoint: {last_checkpoint}\")\n",
    "model_gen.load_weights(last_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YJbQqPFEevop"
   },
   "outputs": [],
   "source": [
    "model_gen.build(tf.TensorShape([1, None]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nQ6JhIGxEARF"
   },
   "outputs": [],
   "source": [
    "model_gen.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Hz3S1j-Ru_Jh"
   },
   "outputs": [],
   "source": [
    "text = \"Why, my dear, you must know, Mrs. Long says that Netherfield is taken by a young man of large fortune from the north of England; that he came down on Monday in a chaise and four to see the place, and was so much delighted with it that he agreed with Mr. Morris immediately; that he is to take possession before Michaelmas, and some of his servants are to be in the house by the end of next week.\"\n",
    "text = text[:params['sequence_len']]\n",
    "while len(text) < params['sequence_len']:\n",
    "    text = ' '+text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "z1osbV7LvNkM"
   },
   "outputs": [],
   "source": [
    "mytext = text[:-1]+'␚'\n",
    "_, gen_text = generate(mytext, model=model_gen, silent=True)\n",
    "print(f\"[{mytext[:-1]}]\",end=\"\")\n",
    "print(gen_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eVbFXeXrriUa"
   },
   "outputs": [],
   "source": [
    "gen_text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WW6LPdlhQtgF"
   },
   "source": [
    "## 6. A dialog with the trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uxDNYZiEQtgF"
   },
   "outputs": [],
   "source": [
    "# Do a dialog with the recursive neural net trained above:\n",
    "# def genDialogAnswer(prompt, g_state=None, endPrompt='.', maxEndPrompts=2,\n",
    "# maxAnswerSize=512, temperature=1.0):\n",
    "\n",
    "\n",
    "def doDialog():\n",
    "    # 0.1 (frozen character) - 1.3 (creative/chaotic character)\n",
    "    temperature = 0.6\n",
    "    endPrompt = '.'  # the endPrompt character is the end-mark in answers.\n",
    "    # look for number of maxEndPrompts until answer is finished.\n",
    "    maxEndPrompts = 4\n",
    "    maxAnswerSize = 2048  # Maximum length of the answer\n",
    "    minAnswerSize = 64  # Minimum length of the answer\n",
    "\n",
    "    print(\"Please enter some dialog.\")\n",
    "    print(\"The net will answer according to your input.\")\n",
    "    print(\"'bye' for end,\")\n",
    "    print(\"'reset' to reset the conversation context,\")\n",
    "    print(\"'temperature=<float>' [0.1(frozen)-1.0(creative)]\")\n",
    "    print(\"    to change character of the dialog.\")\n",
    "    print(\"    Current temperature={}.\".format(temperature))\n",
    "    print()\n",
    "    xso = None\n",
    "    bye = False\n",
    "    doini = True\n",
    "\n",
    "    bye = False\n",
    "    while not bye:\n",
    "        print(\"> \", end=\"\")\n",
    "        prompt = input()\n",
    "        if prompt == 'bye':\n",
    "            bye = True\n",
    "            print(\"Good bye!\")\n",
    "            continue\n",
    "        if prompt == 'reset':\n",
    "            doini = True\n",
    "            print(\"(conversation context marked for reset)\")\n",
    "            continue\n",
    "        if prompt[:len(\"temperature=\")] == \"temperature=\":\n",
    "            t = float(prompt[len(\"temperature=\"):])\n",
    "            if t > 0.05 and t < 1.4:\n",
    "                temperature = t\n",
    "                print(\"(generator temperature now {})\".format(t))\n",
    "                print()\n",
    "                continue\n",
    "            print(\"Invalid temperature-value ignored! [0.1-1.0]\")\n",
    "            continue\n",
    "        _, answer = generate(prompt, model=model_gen, temperature=temperature, gen_len=128, silent=True)\n",
    "        td.source_highlight(answer, min_quote_size=13)\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0JEPK2WIQtgI"
   },
   "outputs": [],
   "source": [
    "# Talk to the net!\n",
    "doDialog()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yWE_ZZMKEARV"
   },
   "source": [
    "## References:\n",
    "* <https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/r2/tutorials/text/text_generation.ipynb>\n",
    "* <https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpu_and_keras.ipynb>"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": "ec3a4d2d-8063-4bfd-a4a2-ee070d3272f7",
   "lastKernelId": "1acc2b74-f51e-477b-910a-a5519dad53b9"
  },
  "accelerator": "TPU",
  "colab": {
   "collapsed_sections": [
    "VmWbteSFQtfq",
    "yWE_ZZMKEARV"
   ],
   "include_colab_link": true,
   "machine_shape": "hm",
   "name": "tensor_poet.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}