generate_model_raw.ipynb

{"cells": [{"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["import os\n", "import time"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["import tensorflow as tf"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["try:\n", "    text = open(\"dataset.txt\", 'rb').read().decode(encoding='utf-8')\n", "except FileNotFoundError:\n", "    raise ValueError(\"can not find dataset\")\n", "print(f'Length of text: {len(text)} characters')\n", "print(text[:250])"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["vocab = sorted(set(text))\n", "print(f'{len(vocab)} unique characters')"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["example_texts = ['abcdefg', 'xyz']\n", "chars = tf.strings.unicode_split(example_texts, input_encoding='UTF-8')\n", "ids_from_chars = tf.keras.layers.StringLookup(\n", "    vocabulary=list(vocab), mask_token=None)\n", "ids = ids_from_chars(chars)\n", "chars_from_ids = tf.keras.layers.StringLookup(\n", "    vocabulary=ids_from_chars.get_vocabulary(), invert=True, mask_token=None)\n", "chars = chars_from_ids(ids)\n", "tf.strings.reduce_join(chars, axis=-1).numpy()"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["def text_from_ids(ids_):\n", "    return tf.strings.reduce_join(chars_from_ids(ids_), axis=-1)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["all_ids = ids_from_chars(tf.strings.unicode_split(text, 'UTF-8'))\n", "print(all_ids)\n", "ids_dataset = tf.data.Dataset.from_tensor_slices(all_ids)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["for ids in ids_dataset.take(10):\n", "    print(chars_from_ids(ids).numpy().decode('utf-8'))"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["seq_length = 100\n", "examples_per_epoch = len(text) // (seq_length + 1)\n", "sequences = ids_dataset.batch(seq_length + 1, drop_remainder=True)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["for seq in sequences.take(1):\n", "    print(chars_from_ids(seq))"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["for seq in sequences.take(5):\n", "    print(text_from_ids(seq).numpy())"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["def split_input_target(sequence):\n", "    input_text = sequence[:-1]\n", "    target_text = sequence[1:]\n", "    return input_text, target_text"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["split_input_target(list(\"Tensorflow\"))\n", "dataset = sequences.map(split_input_target)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["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())"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["BATCH_SIZE = 64\n", "BUFFER_SIZE = 10000"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["dataset = (\n", "    dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)\n", ")"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["vocab_size = len(vocab)\n", "embedding_dim = 256\n", "rnn_units = 1024"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["class MyModel(tf.keras.Model):\n", "    def __init__(self, vocabsize, embeddingdim, rnnunits):\n", "        super().__init__(self)\n", "        self.embedding = tf.keras.layers.Embedding(vocabsize, embeddingdim)\n", "        self.gru = tf.keras.layers.GRU(rnnunits,\n", "                                       return_sequences=True,\n", "                                       return_state=True)\n", "        self.dense = tf.keras.layers.Dense(vocabsize)\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", "        return (x, states) if return_state else x"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["model = MyModel(\n", "    vocabsize=len(ids_from_chars.get_vocabulary()),\n", "    embeddingdim=embedding_dim,\n", "    rnnunits=rnn_units)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["example_batch_predictions = None\n", "input_example_batch = target_example_batch = None"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["for input_example_batch, target_example_batch in dataset.take(1):\n", "    # noinspection PyCallingNonCallable\n", "    example_batch_predictions = model(input_example_batch)\n", "    print(example_batch_predictions.shape, \"# (batch_size, sequence_length, vocab_size)\")"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["model.summary()\n", "sampled_indices = tf.random.categorical(example_batch_predictions[0], num_samples=1)\n", "sampled_indices = tf.squeeze(sampled_indices, axis=-1).numpy()"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["print(\"Input:\\n\", text_from_ids(input_example_batch[0]).numpy())\n", "print()\n", "print(\"Next Char Predictions:\\n\", text_from_ids(sampled_indices).numpy())"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["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()"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["print(\"Prediction shape: \", example_batch_predictions.shape, \" # (batch_size, sequence_length, vocab_size)\")\n", "print(\"Mean loss:        \", mean_loss)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["tf.exp(mean_loss).numpy()\n", "model.compile(optimizer='adam', loss=loss)\n", "checkpoint_dir = './training_checkpoints'\n", "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt_{epoch}\")\n", "checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n", "    filepath=checkpoint_prefix,\n", "    save_weights_only=True)\n", "STEPS = 172  # steps to perform\n", "EPOCHS = 40\n", "history = model.fit(dataset, epochs=EPOCHS, callbacks=[checkpoint_callback], steps_per_epoch=STEPS)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["class OneStep(tf.keras.Model):\n", "    def __init__(self, model_, charsfromids, idsfromchars, temperature=1.0):\n", "        super().__init__()\n", "        self.temperature = temperature\n", "        self.model = model_\n", "        self.chars_from_ids = charsfromids\n", "        self.ids_from_chars = idsfromchars\n", "        skip_ids = self.ids_from_chars(['[UNK]'])[:, None]\n", "        sparse_mask = tf.SparseTensor(\n", "            values=[-float('inf')] * len(skip_ids),\n", "            indices=skip_ids,\n", "            dense_shape=[len(ids_from_chars.get_vocabulary())])\n", "        self.prediction_mask = tf.sparse.to_dense(sparse_mask)\n", "    @tf.function\n", "    def generate_one_step(self, inputs, states_=None):\n", "        input_chars = tf.strings.unicode_split(inputs, 'UTF-8')\n", "        input_ids = self.ids_from_chars(input_chars).to_tensor()\n", "        predicted_logits, states_ = self.model(inputs=input_ids, states_=states_,\n", "                                               return_state=True)\n", "        predicted_logits = predicted_logits[:, -1, :]\n", "        predicted_logits = predicted_logits / self.temperature\n", "        predicted_logits = predicted_logits + self.prediction_mask\n", "        predicted_ids = tf.random.categorical(predicted_logits, num_samples=1)\n", "        predicted_ids = tf.squeeze(predicted_ids, axis=-1)\n", "        predicted_chars = self.chars_from_ids(predicted_ids)\n", "        return predicted_chars, states"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["one_step_model = OneStep(model, chars_from_ids, ids_from_chars)\n", "start = time.time()\n", "states = None\n", "next_char = tf.constant(['seks'])\n", "result = [next_char]"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["for n in range(1000):\n", "    next_char, states = one_step_model.generate_one_step(next_char, states_=states)\n", "    result.append(next_char)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["result = tf.strings.join(result)\n", "end = time.time()"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["print(result[0].numpy().decode('utf-8'), '\\n\\n' + '_' * 80)\n", "print('\\nRun time:', end - start)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["start = time.time()\n", "states = None\n", "next_char = tf.constant(['seks', 'seks', 'seks', 'seks', 'seks'])\n", "result = [next_char]"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["for n in range(1000):\n", "    next_char, states = one_step_model.generate_one_step(next_char, states_=states)\n", "    result.append(next_char)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["result = tf.strings.join(result)\n", "end = time.time()"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["print(result, '\\n\\n' + '_' * 80)\n", "print('\\nRun time:', end - start)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["tf.saved_model.save(one_step_model, 'one_step')\n", "one_step_reloaded = tf.saved_model.load('one_step')\n", "states = None\n", "next_char = tf.constant(['seks'])\n", "result = [next_char]"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["for n in range(100):\n", "    next_char, states = one_step_reloaded.generate_one_step(next_char, states_=states)\n", "    result.append(next_char)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["print(tf.strings.join(result)[0].numpy().decode(\"utf-8\"))"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["class CustomTraining(MyModel):\n", "    @tf.function\n", "    def train_step(self, inputs):\n", "        inputs, labels = inputs\n", "        with tf.GradientTape() as tape:\n", "            # noinspection PyCallingNonCallable\n", "            predictions = self(inputs, training=True)\n", "            loss_ = self.loss(labels, predictions)\n", "        grads = tape.gradient(loss_, model.trainable_variables)\n", "        self.optimizer.apply_gradients(zip(grads, model.trainable_variables))\n", "        return {'loss': loss_}"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["model = CustomTraining(\n", "    vocabsize=len(ids_from_chars.get_vocabulary()),\n", "    embeddingdim=embedding_dim,\n", "    rnnunits=rnn_units)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["model.compile(optimizer=tf.keras.optimizers.Adam(),\n", "              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["model.fit(dataset, epochs=1, steps_per_epoch=STEPS)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["mean = tf.metrics.Mean()\n", "epoch = None"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["for epoch in range(EPOCHS):\n", "    start = time.time()\n", "    mean.reset_states()\n", "    for (batch_n, (inp, target)) in enumerate(dataset):\n", "        if len(dataset) < STEPS:\n", "            raise ValueError(\"dataset is not large enough to do this many steps\")\n", "        logs = model.train_step([inp, target])\n", "        mean.update_state(logs['loss'])\n", "        if batch_n % 50 == 0:\n", "            template = f\"Epoch {epoch + 1} Batch {batch_n} Loss {logs['loss']:.4f}\"\n", "            print(template)\n", "        if batch_n == STEPS:\n", "            break\n", "    if (epoch + 1) % 5 == 0:\n", "        model.save_weights(checkpoint_prefix.format(epoch=epoch))\n", "    print()\n", "    print(f'Epoch {epoch + 1} Loss: {mean.result().numpy():.4f}')\n", "    print(f'Time taken for 1 epoch {time.time() - start:.2f} sec')\n", "    print(\"_\" * 80)"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["model.save_weights(checkpoint_prefix.format(epoch=epoch))"]}], "metadata": {"kernelspec": {"display_name": "Python 3", "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.6.4"}}, "nbformat": 4, "nbformat_minor": 2}