train_and_save_model-20k.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "136413c5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Init Plugin\n",
      "Init Graph Optimizer\n",
      "Init Kernel\n",
      "2.5.0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "import numpy as np\n",
    "import os\n",
    "import time\n",
    "\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "3bf9d7f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length of text: 358262 characters\n",
      "_that***************_this***************_with***************_from***************_your***************_have***************_more***************_will***************_home***************_about**************_page***************_search*************_free*****\n"
     ]
    }
   ],
   "source": [
    "full_text = open('20k.txt', 'rb').read().decode(encoding='utf-8')\n",
    "text = ''\n",
    "seq_length = 20\n",
    "for word in full_text.split():\n",
    "    if len(word) < 4:\n",
    "        continue\n",
    "    else:\n",
    "        word = '_' + word\n",
    "        for i in range(0, seq_length-len(word)):\n",
    "            word = word + '*'\n",
    "    text = text + word\n",
    "\n",
    "print(f'Length of text: {len(text)} characters')\n",
    "print(text[:250])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "be6c5319",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "28 unique characters\n",
      "['*', '_', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']\n"
     ]
    }
   ],
   "source": [
    "vocab = sorted(set(text))\n",
    "print(f'{len(vocab)} unique characters')\n",
    "print(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "80c1312c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO try setting mask_token = 'c' and see what the output looks like\n",
    "\n",
    "ids_from_chars = tf.keras.layers.experimental.preprocessing.StringLookup(\n",
    "    vocabulary=list(vocab), mask_token=None)\n",
    "\n",
    "chars_from_ids = tf.keras.layers.experimental.preprocessing.StringLookup(\n",
    "    vocabulary=ids_from_chars.get_vocabulary(), invert=True, mask_token=None)\n",
    "\n",
    "def text_from_ids(ids):\n",
    "  return tf.strings.reduce_join(chars_from_ids(ids), axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "004730c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(2, shape=(), dtype=int64) _\n",
      "tf.Tensor(22, shape=(), dtype=int64) t\n",
      "tf.Tensor(10, shape=(), dtype=int64) h\n",
      "tf.Tensor(3, shape=(), dtype=int64) a\n",
      "tf.Tensor(22, shape=(), dtype=int64) t\n",
      "tf.Tensor(1, shape=(), dtype=int64) *\n",
      "tf.Tensor(1, shape=(), dtype=int64) *\n",
      "tf.Tensor(1, shape=(), dtype=int64) *\n",
      "tf.Tensor(1, shape=(), dtype=int64) *\n",
      "tf.Tensor(1, shape=(), dtype=int64) *\n"
     ]
    }
   ],
   "source": [
    "all_ids = ids_from_chars(tf.strings.unicode_split(text, 'UTF-8'))\n",
    "ids_dataset = tf.data.Dataset.from_tensor_slices(all_ids)\n",
    "\n",
    "for element in ids_dataset.take(10):\n",
    "    print(element, chars_from_ids(element).numpy().decode('utf-8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "cc8b4060",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([ 2 22 10  3 22  1  1  1  1  1  1  1  1  1  1  1  1  1  1  1], shape=(20,), dtype=int64) tf.Tensor(b'_that***************', shape=(), dtype=string)\n",
      "tf.Tensor([ 2 22 10 11 21  1  1  1  1  1  1  1  1  1  1  1  1  1  1  1], shape=(20,), dtype=int64) tf.Tensor(b'_this***************', shape=(), dtype=string)\n",
      "tf.Tensor([ 2 25 11 22 10  1  1  1  1  1  1  1  1  1  1  1  1  1  1  1], shape=(20,), dtype=int64) tf.Tensor(b'_with***************', shape=(), dtype=string)\n",
      "tf.Tensor([ 2  8 20 17 15  1  1  1  1  1  1  1  1  1  1  1  1  1  1  1], shape=(20,), dtype=int64) tf.Tensor(b'_from***************', shape=(), dtype=string)\n",
      "tf.Tensor([ 2 27 17 23 20  1  1  1  1  1  1  1  1  1  1  1  1  1  1  1], shape=(20,), dtype=int64) tf.Tensor(b'_your***************', shape=(), dtype=string)\n"
     ]
    }
   ],
   "source": [
    "# TODO read karpathi to see how he chooses sequence length and if it's important\n",
    "# maybe longer is better? but we don't want to keep context between words\n",
    "\n",
    "# try truncating and padding after the stop token\n",
    "\n",
    "# try training on only 2 letter words \n",
    "\n",
    "sequences = ids_dataset.batch(seq_length, drop_remainder=True)\n",
    "\n",
    "for seq in sequences.take(5):\n",
    "  print(seq, text_from_ids(seq))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "a4ad559c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(\"args_0:0\", shape=(20,), dtype=int64)\n",
      "Input : b'_that***************'\n",
      "Target: b'that****************'\n"
     ]
    }
   ],
   "source": [
    "def split_input_target(sequence):\n",
    "    print(sequence)\n",
    "    input_text = sequence\n",
    "    target_text = tf.concat([sequence[1:], [1]], 0) # pad with extra *\n",
    "    return input_text, target_text\n",
    "dataset = sequences.map(split_input_target)\n",
    "for input_example, target_example in dataset.take(1):\n",
    "    print(\"Input :\", text_from_ids(input_example).numpy())\n",
    "    print(\"Target:\", text_from_ids(target_example).numpy())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "591d2b87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(array([[ 2, 22, 23, ...,  1,  1,  1],\n",
       "         [ 2,  3,  6, ...,  1,  1,  1],\n",
       "         [ 2,  6,  7, ...,  1,  1,  1],\n",
       "         ...,\n",
       "         [ 2, 15, 17, ...,  1,  1,  1],\n",
       "         [ 2, 14,  3, ...,  1,  1,  1],\n",
       "         [ 2, 21, 13, ...,  1,  1,  1]]),\n",
       "  array([[22, 23, 20, ...,  1,  1,  1],\n",
       "         [ 3,  6, 24, ...,  1,  1,  1],\n",
       "         [ 6,  7,  8, ...,  1,  1,  1],\n",
       "         ...,\n",
       "         [15, 17, 22, ...,  1,  1,  1],\n",
       "         [14,  3, 16, ...,  1,  1,  1],\n",
       "         [21, 13,  3, ...,  1,  1,  1]]))]"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BATCH_SIZE = 64\n",
    "\n",
    "# Buffer size to shuffle the dataset\n",
    "# (TF data is designed to work with possibly infinite sequences,\n",
    "# so it doesn't attempt to shuffle the entire sequence in memory. Instead,\n",
    "# it maintains a buffer in which it shuffles elements).\n",
    "BUFFER_SIZE = 10000\n",
    "\n",
    "dataset = (\n",
    "    dataset\n",
    "    .shuffle(BUFFER_SIZE)\n",
    "    .batch(BATCH_SIZE, drop_remainder=True)\n",
    "    .prefetch(tf.data.experimental.AUTOTUNE))\n",
    "\n",
    "list(dataset.take(1).as_numpy_iterator())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57f82b54",
   "metadata": {},
   "source": [
    "### Define the model\n",
    "We use an RNN with GRUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "c0999b88",
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size = len(vocab)\n",
    "embedding_dim = 256\n",
    "rnn_units = 512\n",
    "\n",
    "class RnnGRUModel(tf.keras.Model):\n",
    "  def __init__(self, vocab_size, embedding_dim, rnn_units):\n",
    "    super().__init__(self)\n",
    "    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n",
    "    self.gru = tf.keras.layers.GRU(rnn_units,\n",
    "                                   return_sequences=True,\n",
    "                                   return_state=True)\n",
    "    self.dense = tf.keras.layers.Dense(vocab_size)\n",
    "    \n",
    "  def call(self, inputs, states=None, return_state=False, training=False):\n",
    "    x = inputs\n",
    "    x = self.embedding(x, training=training)\n",
    "    if states is None:\n",
    "      states = self.gru.get_initial_state(x)\n",
    "    x, states = self.gru(x, initial_state=states, training=training)\n",
    "    x = self.dense(x, training=training)\n",
    "    \n",
    "    if return_state:\n",
    "      return x, states\n",
    "    else:\n",
    "      return x\n",
    "\n",
    "  def get_config(self):\n",
    "    return {\n",
    "        \"embedding\": self.embedding,\n",
    "        \"gru\": self.gru,\n",
    "        \"dense\": self.dense,\n",
    "    }\n",
    "\n",
    "  @classmethod\n",
    "  def from_config(cls, config):\n",
    "    return cls(**config)\n",
    "\n",
    "\n",
    "\n",
    "model = RnnGRUModel(\n",
    "    # Be sure the vocabulary size matches the `StringLookup` layers.\n",
    "    vocab_size=len(ids_from_chars.get_vocabulary()),\n",
    "    embedding_dim=embedding_dim,\n",
    "    rnn_units=rnn_units)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5382190",
   "metadata": {},
   "source": [
    "### Try the untrained model on the first sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "396fea63",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 20, 29) # (batch_size, sequence_length, vocab_size)\n",
      "Model: \"rnn_gru_model\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding (Embedding)        multiple                  7424      \n",
      "_________________________________________________________________\n",
      "gru (GRU)                    multiple                  1182720   \n",
      "_________________________________________________________________\n",
      "dense (Dense)                multiple                  14877     \n",
      "=================================================================\n",
      "Total params: 1,205,021\n",
      "Trainable params: 1,205,021\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "for input_example_batch, target_example_batch in dataset.take(1):\n",
    "    example_batch_predictions = model(input_example_batch)\n",
    "    print(example_batch_predictions.shape, \"# (batch_size, sequence_length, vocab_size)\")\n",
    "    \n",
    "# vocab size is 29 because of [UNK]\n",
    "# TODO we don't actually need [UNK] since input involves all possible characters already\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "5cedc587",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 2, 21,  1, 21, 11,  1,  7, 18, 12, 14,  4,  4, 26,  7, 11, 25, 26,\n",
       "       23, 14, 25])"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sampled_indices = tf.random.categorical(example_batch_predictions[0], num_samples=1)\n",
    "sampled_indices = tf.squeeze(sampled_indices, axis=-1).numpy()\n",
    "sampled_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "aa643e17",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input:\n",
      " _save***************\n",
      "Next Char Predictions:\n",
      " _s*si*epjlbbxeiwxulw\n"
     ]
    }
   ],
   "source": [
    "print(\"Input:\\n\", text_from_ids(input_example_batch[0]).numpy().decode('utf-8'))\n",
    "print(\"Next Char Predictions:\\n\", text_from_ids(sampled_indices).numpy().decode('utf-8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "2891ae18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction shape:  (64, 20, 29)  # (batch_size, sequence_length, vocab_size)\n",
      "Mean loss:         3.3746936\n"
     ]
    }
   ],
   "source": [
    "# Could try with other loss functions, e.g. CosineSimilarity\n",
    "\n",
    "loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "example_batch_loss = loss(target_example_batch, example_batch_predictions)\n",
    "mean_loss = example_batch_loss.numpy().mean()\n",
    "print(\"Prediction shape: \", example_batch_predictions.shape, \" # (batch_size, sequence_length, vocab_size)\")\n",
    "print(\"Mean loss:        \", mean_loss)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db7f8784",
   "metadata": {},
   "source": [
    "### Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "9f2ce834",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(optimizer='adam', loss=loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "a189390c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Directory where the checkpoints will be saved\n",
    "checkpoint_dir = '20k_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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "36afa12a",
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "e41138f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/30\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-01-17 15:16:27.899779: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:112] Plugin optimizer for device_type GPU is enabled.\n",
      "2022-01-17 15:16:28.106708: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:112] Plugin optimizer for device_type GPU is enabled.\n",
      "2022-01-17 15:16:28.771316: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:112] Plugin optimizer for device_type GPU is enabled.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "279/279 [==============================] - 14s 41ms/step - loss: 1.2215\n",
      "Epoch 2/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 1.0584\n",
      "Epoch 3/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.9947\n",
      "Epoch 4/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.9366\n",
      "Epoch 5/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.9019\n",
      "Epoch 6/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.8732\n",
      "Epoch 7/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.8580\n",
      "Epoch 8/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.8310\n",
      "Epoch 9/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.8037\n",
      "Epoch 10/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.7867\n",
      "Epoch 11/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.7686\n",
      "Epoch 12/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.7517\n",
      "Epoch 13/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.7348\n",
      "Epoch 14/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.7169\n",
      "Epoch 15/30\n",
      "279/279 [==============================] - 12s 41ms/step - loss: 0.7029\n",
      "Epoch 16/30\n",
      "279/279 [==============================] - 12s 41ms/step - loss: 0.6897\n",
      "Epoch 17/30\n",
      "279/279 [==============================] - 12s 41ms/step - loss: 0.6795\n",
      "Epoch 18/30\n",
      "279/279 [==============================] - 12s 41ms/step - loss: 0.6654\n",
      "Epoch 19/30\n",
      "279/279 [==============================] - 12s 42ms/step - loss: 0.6534\n",
      "Epoch 20/30\n",
      "279/279 [==============================] - 12s 42ms/step - loss: 0.6435\n",
      "Epoch 21/30\n",
      "279/279 [==============================] - 12s 41ms/step - loss: 0.6363\n",
      "Epoch 22/30\n",
      "279/279 [==============================] - 12s 41ms/step - loss: 0.6308\n",
      "Epoch 23/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.6272\n",
      "Epoch 24/30\n",
      "279/279 [==============================] - 11s 39ms/step - loss: 0.6209\n",
      "Epoch 25/30\n",
      "279/279 [==============================] - 11s 39ms/step - loss: 0.6132\n",
      "Epoch 26/30\n",
      "279/279 [==============================] - 11s 39ms/step - loss: 0.6076\n",
      "Epoch 27/30\n",
      "279/279 [==============================] - 11s 38ms/step - loss: 0.6050\n",
      "Epoch 28/30\n",
      "279/279 [==============================] - 11s 38ms/step - loss: 0.6071\n",
      "Epoch 29/30\n",
      "279/279 [==============================] - 11s 40ms/step - loss: 0.6087\n",
      "Epoch 30/30\n",
      "279/279 [==============================] - 12s 41ms/step - loss: 0.6024\n"
     ]
    }
   ],
   "source": [
    "# TODO: add a callback to print out generated words at a point in training\n",
    "# TODO: add a callback to print out how the model evaluates a particular word at each point in training\n",
    "\n",
    "history = model.fit(dataset, epochs=EPOCHS, callbacks=[checkpoint_callback])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a6deae8",
   "metadata": {},
   "source": [
    "### Create ModelWrapper which samples from the model one step at a time and evaluates words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "7d02212c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModelWrapper():\n",
    "  def __init__(self, model, chars_from_ids, ids_from_chars, temperature=1):\n",
    "    self.temperature = temperature\n",
    "    self.model = model\n",
    "    self.chars_from_ids = chars_from_ids\n",
    "    self.ids_from_chars = ids_from_chars\n",
    "\n",
    "  @tf.function\n",
    "  def generate_one_step(self, inputs, states=None):\n",
    "    # Convert strings to token IDs.\n",
    "    input_chars = tf.strings.unicode_split(inputs, 'UTF-8')\n",
    "    input_ids = self.ids_from_chars(input_chars).to_tensor()\n",
    "\n",
    "    # Run the model.\n",
    "    # predicted_logits.shape is [batch, char, next_char_logits]\n",
    "    predicted_logits, states = self.model(inputs=input_ids, states=states,\n",
    "                                          return_state=True)\n",
    "    # Only use the last prediction.\n",
    "    predicted_logits = predicted_logits[:, -1, :]\n",
    "    predicted_logits = predicted_logits/self.temperature\n",
    "\n",
    "    # Sample the output logits to generate token IDs.\n",
    "    predicted_ids = tf.random.categorical(predicted_logits, num_samples=1)\n",
    "    predicted_ids = tf.squeeze(predicted_ids, axis=-1)\n",
    "\n",
    "    # Convert from token ids to characters\n",
    "    predicted_chars = self.chars_from_ids(predicted_ids)\n",
    "\n",
    "    # Return the characters and model state.\n",
    "    return predicted_chars, states\n",
    "\n",
    "  def probability_of_letter(this, logits, letter):\n",
    "    return (tf.squeeze(logits)[this.ids_from_chars(letter).numpy()]).numpy()\n",
    "\n",
    "  def evaluate_word(this, word, show_work):\n",
    "      # Seed the model with '**'\n",
    "      seed_chars = tf.constant(['**'])\n",
    "      input_chars = tf.strings.unicode_split(seed_chars, 'UTF-8')\n",
    "      input_ids = this.ids_from_chars(input_chars).to_tensor()\n",
    "      states = None\n",
    "\n",
    "      # get the initial probability distribution\n",
    "      predicted_logits, states = this.model(inputs=input_ids, states=states, return_state=True)\n",
    "      predicted_logits = predicted_logits[:, -1, :]/this.temperature\n",
    "\n",
    "      total_prob = 0\n",
    "      for letter in word:\n",
    "        # check the log probability of the letter given current state\n",
    "        probability = this.probability_of_letter(predicted_logits, letter)\n",
    "        total_prob += probability\n",
    "        if show_work:\n",
    "          print(\"Probability of\", letter, \"=\", probability)\n",
    "\n",
    "        # feed this letter into the model\n",
    "        this_char = tf.constant([letter])        \n",
    "        input_char = tf.strings.unicode_split(this_char, 'UTF-8')\n",
    "        this_input_id = this.ids_from_chars(input_char).to_tensor()\n",
    "        predicted_logits, states = this.model(inputs=this_input_id, states=states,\n",
    "                                              return_state=True)\n",
    "        predicted_logits = predicted_logits[:, -1, :]/this.temperature\n",
    "\n",
    "      # check the log probability of '*' given current state\n",
    "      probability = this.probability_of_letter(predicted_logits, '*')\n",
    "      # TODO This assumes that the probability of the word is the product of the probability of each letter\n",
    "      # the problem is that I'm not really training the RNN to generate probabilities for each letter. In otherwords\n",
    "      # the log likelihoods in the logits is not real, just directionally correct\n",
    "      total_prob += probability\n",
    "      if show_work:\n",
    "        print(\"Probability of * =\", probability)\n",
    "\n",
    "      # TODO normalize by length\n",
    "      # TODO actually, shouldn't the model do this automatically if it's well trained?\n",
    "      return total_prob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "c607d4f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_wrapper = ModelWrapper(model, chars_from_ids, ids_from_chars, temperature=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "27958a31",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_words(given_model, num_words):\n",
    "    for m in range(num_words):\n",
    "      states = None\n",
    "      next_char = tf.constant(['_'])\n",
    "      result = [next_char]\n",
    "\n",
    "      for n in range(100):\n",
    "        next_char, states = given_model.generate_one_step(next_char, states=states)\n",
    "        result.append(next_char)\n",
    "        if next_char == '*':\n",
    "          break\n",
    "\n",
    "      result = tf.strings.join(result)\n",
    "      end = time.time()\n",
    "      print(result[0].numpy().decode('utf-8'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "07e5104c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_workings*\n",
      "_bouldigh*\n",
      "_flyer*\n",
      "_completely*\n",
      "_stored*\n",
      "_disable*\n",
      "_women*\n",
      "_summerie*\n",
      "_alterers*\n",
      "_downlodk*\n",
      "_bioteerahion*\n",
      "_salled*\n",
      "_toren*\n",
      "_importing*\n",
      "_hurfy*\n",
      "_viewers*\n",
      "_durby*\n",
      "_fins*\n",
      "_authorizer*\n",
      "_bulldence*\n",
      "_rocks*\n",
      "_link*\n",
      "_sundancing*\n",
      "_tapfued*\n",
      "_strippers*\n",
      "_inspected*\n",
      "_chimpionship*\n",
      "_copand*\n",
      "_fulfillment*\n",
      "_turh*\n",
      "_selections*\n",
      "_roommate*\n",
      "_drawa*\n",
      "_buddsets*\n",
      "_ktrides*\n",
      "_unpeaut*\n",
      "_physiologiss*\n",
      "_harder*\n",
      "_alphobe*\n",
      "_leaker*\n",
      "_faeron*\n",
      "_christi*\n",
      "_femoding*\n",
      "_fayle*\n",
      "_browson*\n",
      "_nozhare*\n",
      "_communisment*\n",
      "_horizon*\n",
      "_mankigh*\n",
      "_gift*\n"
     ]
    }
   ],
   "source": [
    "create_words(model_wrapper, 50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2150534",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tensorflow25] *",
   "language": "python",
   "name": "conda-env-tensorflow25-py"
  },
  "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}