Transformer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Transformer",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyP+nOVev9fA/wmj4Cxgm74Z"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "hHRo2NsK-Rgd",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Import & seed\n",
        "import time\n",
        "import re\n",
        "import os\n",
        "import math\n",
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "import requests\n",
        "import collections\n",
        "import pickle\n",
        "import copy, random\n",
        "import nltk as nl\n",
        "nl.download('punkt')\n",
        "from itertools import zip_longest\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras.models import Sequential, Model\n",
        "from tensorflow.keras.layers import (\n",
        "    Input, Reshape, BatchNormalization, Dense, Dropout, concatenate,\n",
        "    Embedding, LSTM, Dense, GRU, Bidirectional, Add\n",
        ")\n",
        "from tensorflow.keras.activations import elu, relu, softmax, sigmoid\n",
        "from tensorflow.keras import regularizers\n",
        "from tensorflow.keras.utils import to_categorical\n",
        "\n",
        "print(tf.__version__)\n",
        "\n",
        "np.random.seed(1234)\n",
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qLZ0LreHzhhr",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Setup wandb\n",
        "!pip install wandb\n",
        "!wandb login #######################################\n",
        "import wandb\n",
        "wandb.init(project=\"deep_comedy\", name=\"optimal\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "onTHkWW_YOHp",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Model\n",
        "\n",
        "vocab_size = 1800\n",
        "terces_per_batch = 4\n",
        "terces_len = 75\n",
        "batch_len = terces_per_batch * (terces_len + 1)\n",
        "\n",
        "wandb.config.num_layers = 4\n",
        "wandb.config.d_model = 128\n",
        "wandb.config.dff = 256\n",
        "wandb.config.num_heads = 4\n",
        "wandb.config.dropout = 0.1\n",
        "wandb.config.learning_rate = 2e-4 \n",
        "\n",
        "EPOCHS = 150\n",
        "generate_at = []\n",
        "\n",
        "def get_angles(pos, i, d_model):\n",
        "    angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n",
        "    return pos * angle_rates\n",
        "\n",
        "def positional_encoding(position, d_model):\n",
        "    angle_rads = get_angles(np.arange(position)[:, np.newaxis],\n",
        "                            np.arange(d_model)[np.newaxis, :],\n",
        "                            d_model)\n",
        "    # apply sin to even indices in the array; 2i\n",
        "    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
        "    # apply cos to odd indices in the array; 2i+1\n",
        "    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
        "    pos_encoding = angle_rads[np.newaxis, ...]\n",
        "    return tf.cast(pos_encoding, dtype=tf.float32)\n",
        "\n",
        "def create_padding_mask(seq):   \n",
        "    seq = tf.cast(tf.math.equal(seq, pad), tf.float32)\n",
        "    # add extra dimensions to add the padding\n",
        "    # to the attention logits.\n",
        "    return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)\n",
        "\n",
        "def create_look_ahead_mask(size):\n",
        "    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n",
        "    return mask  # (seq_len, seq_len)\n",
        "\n",
        "def scaled_dot_product_attention(q, k, v, mask):\n",
        "    \"\"\"Calculate the attention weights.\n",
        "    q, k, v must have matching leading dimensions.\n",
        "    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.\n",
        "    The mask has different shapes depending on its type(padding or look ahead) \n",
        "    but it must be broadcastable for addition.\n",
        "    \n",
        "    Args:\n",
        "        q: query shape == (..., seq_len_q, depth)\n",
        "        k: key shape == (..., seq_len_k, depth)\n",
        "        v: value shape == (..., seq_len_v, depth_v)\n",
        "        mask: Float tensor with shape broadcastable \n",
        "            to (..., seq_len_q, seq_len_k). Defaults to None.\n",
        "        \n",
        "    Returns:\n",
        "        output, attention_weights\n",
        "    \"\"\"\n",
        "    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)\n",
        "    # scale matmul_qk\n",
        "    dk = tf.cast(tf.shape(k)[-1], tf.float32)\n",
        "    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)\n",
        "    # add the mask to the scaled tensor.\n",
        "    if mask is not None:\n",
        "        scaled_attention_logits += (mask * -1e9)  \n",
        "    # softmax is normalized on the last axis (seq_len_k) so that the scores\n",
        "    # add up to 1.\n",
        "    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)\n",
        "    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)\n",
        "    return output, attention_weights\n",
        "\n",
        "def point_wise_feed_forward_network(d_model, dff):\n",
        "    return tf.keras.Sequential([\n",
        "        tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)\n",
        "        tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)\n",
        "    ])\n",
        "\n",
        "\n",
        "class MultiHeadAttention(tf.keras.layers.Layer):\n",
        "    def __init__(self, d_model, num_heads):\n",
        "        super(MultiHeadAttention, self).__init__()\n",
        "        self.num_heads = num_heads\n",
        "        self.d_model = d_model\n",
        "        assert d_model % self.num_heads == 0\n",
        "        self.depth = d_model // self.num_heads\n",
        "        self.wq = tf.keras.layers.Dense(d_model)\n",
        "        self.wk = tf.keras.layers.Dense(d_model)\n",
        "        self.wv = tf.keras.layers.Dense(d_model)\n",
        "        self.dense = tf.keras.layers.Dense(d_model)\n",
        "            \n",
        "    def split_heads(self, x, batch_size):\n",
        "        \"\"\"Split the last dimension into (num_heads, depth).\n",
        "        Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)\n",
        "        \"\"\"\n",
        "        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n",
        "        return tf.transpose(x, perm=[0, 2, 1, 3])\n",
        "        \n",
        "    def call(self, v, k, q, mask):\n",
        "        batch_size = tf.shape(q)[0]\n",
        "        \n",
        "        q = self.wq(q)  # (batch_size, seq_len, d_model)\n",
        "        k = self.wk(k)  # (batch_size, seq_len, d_model)\n",
        "        v = self.wv(v)  # (batch_size, seq_len, d_model)\n",
        "        \n",
        "        q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)\n",
        "        k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)\n",
        "        v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)\n",
        "        \n",
        "        # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)\n",
        "        # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)\n",
        "        scaled_attention, attention_weights = scaled_dot_product_attention(q, k, v, mask)\n",
        "        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)\n",
        "        concat_attention = tf.reshape(scaled_attention, \n",
        "                                    (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)\n",
        "        output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)\n",
        "        return output, attention_weights\n",
        "\n",
        "class EncoderLayer(tf.keras.layers.Layer):\n",
        "    def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
        "        super(EncoderLayer, self).__init__()\n",
        "\n",
        "        self.mha = MultiHeadAttention(d_model, num_heads)\n",
        "        self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
        "\n",
        "        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        \n",
        "        self.dropout1 = tf.keras.layers.Dropout(rate)\n",
        "        self.dropout2 = tf.keras.layers.Dropout(rate)\n",
        "        \n",
        "    def __call__(self, x, training, mask):\n",
        "        attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)\n",
        "        attn_output = self.dropout1(attn_output, training=training)\n",
        "        out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)\n",
        "\n",
        "        ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)\n",
        "        ffn_output = self.dropout2(ffn_output, training=training)\n",
        "        out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)\n",
        "        \n",
        "        return out2\n",
        "\n",
        "class DecoderLayer(tf.keras.layers.Layer):\n",
        "    def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
        "        super(DecoderLayer, self).__init__()\n",
        "\n",
        "        self.mha1 = MultiHeadAttention(d_model, num_heads)\n",
        "        self.mha2 = MultiHeadAttention(d_model, num_heads)\n",
        "\n",
        "        self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
        "    \n",
        "        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        \n",
        "        self.dropout1 = tf.keras.layers.Dropout(rate)\n",
        "        self.dropout2 = tf.keras.layers.Dropout(rate)\n",
        "        self.dropout3 = tf.keras.layers.Dropout(rate)\n",
        "        \n",
        "        \n",
        "    def __call__(self, x, enc_output, training, \n",
        "            look_ahead_mask, padding_mask):\n",
        "        # enc_output.shape == (batch_size, input_seq_len, d_model)\n",
        "\n",
        "        attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)\n",
        "        attn1 = self.dropout1(attn1, training=training)\n",
        "        out1 = self.layernorm1(attn1 + x)\n",
        "        \n",
        "        attn2, attn_weights_block2 = self.mha2(\n",
        "            enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)\n",
        "        attn2 = self.dropout2(attn2, training=training)\n",
        "        out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)\n",
        "        \n",
        "        ffn_output = self.ffn(out2)  # (batch_size, input_seq_len, d_model)\n",
        "        ffn_output = self.dropout3(ffn_output, training=training)\n",
        "        out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)\n",
        "        \n",
        "        return out3, attn_weights_block1, attn_weights_block2\n",
        "\n",
        "class Encoder(tf.keras.layers.Layer):\n",
        "    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,\n",
        "                maximum_position_encoding, rate=0.1):\n",
        "        super(Encoder, self).__init__()\n",
        "        self.d_model = d_model\n",
        "        self.num_layers = num_layers\n",
        "        self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)\n",
        "        self.pos_encoding = positional_encoding(maximum_position_encoding, \n",
        "                                                self.d_model)\n",
        "        self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) \n",
        "                        for _ in range(num_layers)]\n",
        "        self.dropout = tf.keras.layers.Dropout(rate)\n",
        "            \n",
        "    def __call__(self, x, training, mask):\n",
        "        seq_len = tf.shape(x)[1]\n",
        "        # adding embedding and position encoding.\n",
        "        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)\n",
        "        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
        "        x += self.pos_encoding[:, :seq_len, :]\n",
        "        x = self.dropout(x, training=training)\n",
        "        for i in range(self.num_layers):\n",
        "            x = self.enc_layers[i](x, training, mask)\n",
        "        return x  # (batch_size, input_seq_len, d_model)\n",
        "\n",
        "class Decoder(tf.keras.layers.Layer):\n",
        "    def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,\n",
        "                maximum_position_encoding, rate=0.1):\n",
        "        super(Decoder, self).__init__()\n",
        "        self.d_model = d_model\n",
        "        self.num_layers = num_layers\n",
        "        self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)\n",
        "        self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)\n",
        "        self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate) \n",
        "                        for _ in range(num_layers)]\n",
        "        self.dropout = tf.keras.layers.Dropout(rate)\n",
        "        \n",
        "    def __call__(self, x, enc_output, training, look_ahead_mask, padding_mask):\n",
        "        seq_len = tf.shape(x)[1]\n",
        "        attention_weights = {}\n",
        "        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)\n",
        "        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
        "        x += self.pos_encoding[:, :seq_len, :]\n",
        "        x = self.dropout(x, training=training)\n",
        "        for i in range(self.num_layers):\n",
        "            x, block1, block2 = self.dec_layers[i](x, enc_output, training,\n",
        "                                                look_ahead_mask, padding_mask)\n",
        "            attention_weights['decoder_layer{}_block1'.format(i+1)] = block1\n",
        "            attention_weights['decoder_layer{}_block2'.format(i+1)] = block2\n",
        "        \n",
        "        # x.shape == (batch_size, target_seq_len, d_model)\n",
        "        return x, attention_weights\n",
        "\n",
        "class Transformer(tf.keras.Model):\n",
        "    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, \n",
        "                target_vocab_size, pe_input, pe_target, rate=0.1):\n",
        "        super(Transformer, self).__init__()\n",
        "\n",
        "        self.encoder = Encoder(num_layers, d_model, num_heads, dff, \n",
        "                            input_vocab_size, pe_input, rate)\n",
        "\n",
        "        self.decoder = Decoder(num_layers, d_model, num_heads, dff, \n",
        "                            target_vocab_size, pe_target, rate)\n",
        "\n",
        "        self.final_layer = tf.keras.layers.Dense(target_vocab_size)\n",
        "        \n",
        "    def __call__(self, inp, tar, training, enc_padding_mask, \n",
        "            look_ahead_mask, dec_padding_mask):\n",
        "\n",
        "        enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)\n",
        "        \n",
        "        # dec_output.shape == (batch_size, tar_seq_len, d_model)\n",
        "        dec_output, attention_weights = self.decoder(\n",
        "            tar, enc_output, training, look_ahead_mask, dec_padding_mask)\n",
        "        \n",
        "        final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)\n",
        "        \n",
        "        return final_output, attention_weights\n",
        "\n",
        "optimizer = tf.keras.optimizers.Adam(wandb.config.learning_rate, beta_1=0.9, beta_2=0.98, \n",
        "                                     epsilon=1e-9)\n",
        "\n",
        "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True, reduction='none')\n",
        "\n",
        "def loss_function(real, pred):\n",
        "    mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
        "    loss_ = loss_object(real, pred)\n",
        "    mask = tf.cast(mask, dtype=loss_.dtype)\n",
        "    loss_ *= mask\n",
        "    return tf.reduce_sum(loss_)/tf.reduce_sum(mask)\n",
        "\n",
        "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
        "train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
        "    name='train_accuracy')\n",
        "\n",
        "val_loss = tf.keras.metrics.Mean(name='train_loss')\n",
        "val_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
        "    name='train_accuracy')\n",
        "\n",
        "input_vocab_size = vocab_size\n",
        "target_vocab_size = vocab_size\n",
        "max_len = batch_len\n",
        "\n",
        "transformer = Transformer(wandb.config.num_layers, wandb.config.d_model, \n",
        "                          wandb.config.num_heads, wandb.config.dff,\n",
        "                          input_vocab_size, target_vocab_size, \n",
        "                          pe_input=max_len, \n",
        "                          pe_target=max_len,\n",
        "                          rate=wandb.config.dropout)\n",
        "def create_masks(inp, tar):\n",
        "    # Encoder padding mask\n",
        "    enc_padding_mask = create_padding_mask(inp)\n",
        "    \n",
        "    # Used in the 2nd attention block in the decoder.\n",
        "    # This padding mask is used to mask the encoder outputs.\n",
        "    dec_padding_mask = create_padding_mask(inp)\n",
        "    \n",
        "    # Used in the 1st attention block in the decoder.\n",
        "    # It is used to pad and mask future tokens in the input received by \n",
        "    # the decoder.\n",
        "    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])\n",
        "    dec_target_padding_mask = create_padding_mask(tar)\n",
        "    combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n",
        "    \n",
        "    return enc_padding_mask, combined_mask, dec_padding_mask\n",
        "\n",
        "checkpoint_path = \"./checkpoints/train\"\n",
        "\n",
        "ckpt = tf.train.Checkpoint(transformer=transformer,\n",
        "                           optimizer=optimizer)\n",
        "ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)\n",
        "# if a checkpoint exists, restore the latest checkpoint.\n",
        "if ckpt_manager.latest_checkpoint:\n",
        "    ckpt.restore(ckpt_manager.latest_checkpoint)\n",
        "    print ('Latest checkpoint restored!!')\n",
        "\n",
        "@tf.function()\n",
        "def train_step(inp, tar):\n",
        "    tar_inp = tar[:, :-1]\n",
        "    tar_real = tar[:, 1:]\n",
        "    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)\n",
        "    \n",
        "    with tf.GradientTape() as tape:\n",
        "        predictions, _ = transformer(inp, tar_inp, \n",
        "                                    True, \n",
        "                                    enc_padding_mask, \n",
        "                                    combined_mask, \n",
        "                                    dec_padding_mask)\n",
        "        loss = loss_function(tar_real, predictions)\n",
        "        gradients = tape.gradient(loss, transformer.trainable_variables)    \n",
        "        optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))\n",
        "    \n",
        "    train_loss(loss)\n",
        "    train_accuracy(tar_real, predictions)\n",
        "\n",
        "@tf.function()\n",
        "def val_step(inp, tar):\n",
        "    tar_inp = tar[:, :-1]\n",
        "    tar_real = tar[:, 1:]\n",
        "    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)\n",
        "    \n",
        "    predictions, _ = transformer(inp, tar_inp, \n",
        "                                False, \n",
        "                                enc_padding_mask, \n",
        "                                combined_mask, \n",
        "                                dec_padding_mask)\n",
        "    loss = loss_function(tar_real, predictions)\n",
        "    \n",
        "    val_loss(loss)\n",
        "    val_accuracy(tar_real, predictions)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GI14h2PZX70w",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Preprocessing\n",
        "\n",
        "def get_hyp_lm_tercets(tercets):\n",
        "    new_tercets = []\n",
        "    for tercet in tercets:\n",
        "        new_tercets.append([])\n",
        "        for verse in tercet:\n",
        "            new_tercets[-1].append([])\n",
        "            for hyp_w in verse:\n",
        "                new_tercets[-1][-1].extend(hyp_w)\n",
        "                new_tercets[-1][-1].append('<SEP>')\n",
        "            new_tercets[-1][-1] = new_tercets[-1][-1][:-1]\n",
        "\n",
        "    return new_tercets\n",
        "\n",
        "def is_vowel(c):\n",
        "    return c in 'aeiouAEIOUàìíèéùúüòï'\n",
        "\n",
        "def unsplittable_cons():\n",
        "    u_cons = []\n",
        "    for c1 in ('b', 'c', 'd', 'f', 'g', 'p', 't', 'v'):\n",
        "        for c2 in ('l', 'r'):\n",
        "            u_cons.append(c1 + c2)\n",
        "\n",
        "    others = ['gn', 'gh', 'ch']\n",
        "    u_cons.extend(others)\n",
        "    return u_cons\n",
        "\n",
        "\n",
        "def are_cons_to_split(c1, c2):\n",
        "    to_split = ('cq', 'cn', 'lm', 'rc', 'bd', 'mb', 'mn', 'ld', 'ng', 'nd', 'tm', 'nv', 'nc', 'ft', 'nf', 'gm', 'fm', 'rv', 'fp')\n",
        "    return (c1 + c2) in to_split or (not is_vowel(c1) and (c1 == c2)) or ((c1 + c2) not in unsplittable_cons()) and (\n",
        "        (not is_vowel(c1)) and (not is_vowel(c2)) and c1 != 's')\n",
        "\n",
        "\n",
        "def is_diphthong(c1, c2):\n",
        "    return (c1 + c2) in ('ia', 'ie', 'io', 'iu', 'ua', 'ue', 'uo', 'ui', 'ai', 'ei', 'oi', 'ui', 'au', 'eu', 'ïe', 'iú', 'iù')\n",
        "\n",
        "\n",
        "def is_triphthong(c1, c2, c3):\n",
        "    return (c1 + c2 + c3) in ('iai', 'iei', 'uoi', 'uai', 'uei', 'iuo')\n",
        "\n",
        "\n",
        "def is_toned_vowel(c):\n",
        "    return c in 'àìèéùòï'\n",
        "\n",
        "def has_vowels(sy):\n",
        "    for c in sy:\n",
        "        if is_vowel(c):\n",
        "            return True\n",
        "    return False\n",
        "\n",
        "\n",
        "def hyphenation(word):\n",
        "    \"\"\"\n",
        "    Split word in syllables\n",
        "    :param word: input string\n",
        "    :return: a list containing syllables of the word\n",
        "    \"\"\"\n",
        "    if not word or word == '':\n",
        "        return []\n",
        "    # elif len(word) == 3 and (is_vowel(word[1]) and is_vowel(word[2]) and not is_toned_vowel(word[2]) and (\n",
        "    #     not is_diphthong(word[1], word[2]) or (word[1] == 'i'))):\n",
        "    elif len(word) == 3 and (is_vowel(word[1]) and is_vowel(word[2]) and not is_toned_vowel(word[2]) and (\n",
        "        not is_diphthong(word[1], word[2]))):\n",
        "        return [word[:2]] + [word[2]]\n",
        "    elif len(word) == 3 and is_vowel(word[0]) and not is_vowel(word[1]) and is_vowel(word[2]):\n",
        "        return [word[:2]] + [word[2]]\n",
        "    elif len(word) == 3:\n",
        "        return [word]\n",
        "\n",
        "    syllables = []\n",
        "    is_done = False\n",
        "    count = 0\n",
        "    while not is_done and count <= len(word) - 1:\n",
        "        syllables.append('')\n",
        "        c = word[count]\n",
        "        while not is_vowel(c) and count < len(word) - 1:\n",
        "            syllables[-1] = syllables[-1] + c\n",
        "            count += 1\n",
        "            c = word[count]\n",
        "\n",
        "        syllables[-1] = syllables[-1] + word[count]\n",
        "\n",
        "        if count == len(word) - 1:\n",
        "            is_done = True\n",
        "        else:\n",
        "            count += 1\n",
        "\n",
        "            if count < len(word) and not is_vowel(word[count]):\n",
        "                if count == len(word) - 1:\n",
        "                    syllables[-1] += word[count]\n",
        "                    count += 1\n",
        "                elif count + 1 < len(word) and are_cons_to_split(word[count], word[count + 1]):\n",
        "                    syllables[-1] += word[count]\n",
        "                    count += 1\n",
        "                elif count + 2 < len(word) and not is_vowel(word[count + 1]) and not is_vowel(word[count + 2]) and word[\n",
        "                    count] != 's':\n",
        "                    syllables[-1] += word[count]\n",
        "                    count += 1\n",
        "            elif count < len(word):\n",
        "                if count + 1 < len(word) and is_triphthong(word[count - 1], word[count], word[count + 1]):\n",
        "                    syllables[-1] += word[count] + word[count + 1]\n",
        "                    count += 2\n",
        "                elif is_diphthong(word[count - 1], word[count]):\n",
        "                    syllables[-1] += word[count]\n",
        "                    count += 1\n",
        "\n",
        "                if count + 1 < len(word) and are_cons_to_split(word[count], word[count + 1]):\n",
        "                    syllables[-1] += word[count]\n",
        "                    count += 1\n",
        "\n",
        "            else:\n",
        "                is_done = True\n",
        "\n",
        "    if not has_vowels(syllables[-1]) and len(syllables) > 1:\n",
        "        syllables[-2] = syllables[-2] + syllables[-1]\n",
        "        syllables = syllables[:-1]\n",
        "\n",
        "    return syllables\n",
        "\n",
        "\n",
        "\n",
        "def get_dc_hyphenation(canti):\n",
        "    hyp_canti, hyp_tokens = [], []\n",
        "    for canto in canti:\n",
        "        hyp_canti.append([])\n",
        "        for verso in canto:\n",
        "            syllables = seq_hyphentation(verso)\n",
        "            hyp_canti[-1].append(syllables)\n",
        "            for syllable in syllables:\n",
        "                hyp_tokens.extend(syllable)\n",
        "\n",
        "    return hyp_canti, hyp_tokens\n",
        "\n",
        "\n",
        "def seq_hyphentation(words):\n",
        "    \"\"\"\n",
        "    Converts words in a list of strings into lists of syllables\n",
        "    :param words: a list of words (strings)\n",
        "    :return: a list of lists containing word syllables\n",
        "    \"\"\"\n",
        "    return [hyphenation(w) for w in words]\n",
        "\n",
        "\n",
        "def get_dc_cantos(filename, encoding=None):\n",
        "    # raw_data = read_words(filename=filename)\n",
        "    cantos, words, raw = [], [], []\n",
        "    with open(filename, \"r\", encoding=encoding) as f:\n",
        "        for line in f:\n",
        "            sentence = line.strip()\n",
        "            sentence = str.replace(sentence, \"\\.\", \" \\. \")\n",
        "            sentence = str.replace(sentence, \"[\", '')\n",
        "            sentence = str.replace(sentence, \"]\", '')\n",
        "            sentence = str.replace(sentence, \"-\", '')\n",
        "            sentence = str.replace(sentence, \";\", \" ; \")\n",
        "            sentence = str.replace(sentence, \",\", \" , \")\n",
        "            # sentence = str.replace(sentence, \" \\'\", '')\n",
        "            sentence = str.replace(sentence, \"\\'\", ' \\' ')\n",
        "            if len(sentence) > 1:\n",
        "                # sentence = sentence.translate(string.punctuation)\n",
        "                tokenized_sentence = nl.word_tokenize(sentence)\n",
        "                # tokenized_sentence = sentence.split()\n",
        "                tokenized_sentence = [w.lower() for w in tokenized_sentence if len(w) > 0]\n",
        "\n",
        "                tokenized_sentence = [w for w in tokenized_sentence if \",\" not in w]\n",
        "                tokenized_sentence = [w for w in tokenized_sentence if \".\" not in w]\n",
        "                tokenized_sentence = [w for w in tokenized_sentence if \":\" not in w]\n",
        "                tokenized_sentence = [w for w in tokenized_sentence if \";\" not in w]\n",
        "                tokenized_sentence = [w for w in tokenized_sentence if \"«\" not in w]\n",
        "                tokenized_sentence = [w for w in tokenized_sentence if \"»\" not in w]\n",
        "                # ts = []\n",
        "                ts = tokenized_sentence\n",
        "                # [ts.extend(re.split(\"(\\')\", e)) for e in tokenized_sentence]\n",
        "                tokenized_sentence = [w for w in ts if len(w) > 0]\n",
        "\n",
        "                if len(tokenized_sentence) == 2:\n",
        "                    cantos.append([])\n",
        "                    raw.append([])\n",
        "                elif len(tokenized_sentence) > 2:\n",
        "                    raw[-1].append(sentence)\n",
        "                    cantos[-1].append(tokenized_sentence)\n",
        "                    words.extend(tokenized_sentence)\n",
        "\n",
        "    return cantos, words, raw\n",
        "\n",
        "\n",
        "def create_tercets(cantos):\n",
        "    tercets = []\n",
        "    for i,canto in enumerate(cantos):\n",
        "        for v,verse in enumerate(canto):\n",
        "            if v%3 == 0:\n",
        "                tercets.append([])\n",
        "\n",
        "            tercets[-1].append(verse)\n",
        "        tercets = tercets[:-1]  # removes the last malformed tercets (only 2 verses)\n",
        "\n",
        "    return tercets\n",
        "\n",
        "def pad_list(l, pad_token, max_l_size, keep_lasts=False, pad_right=True):\n",
        "    \"\"\"\n",
        "    Adds a padding token to a list\n",
        "    inputs:\n",
        "    :param l: input list to pad.\n",
        "    :param pad_token: value to add as padding.\n",
        "    :param max_l_size: length of the new padded list to return,\n",
        "    it truncates lists longer that 'max_l_size' without adding\n",
        "    padding values.\n",
        "    :param keep_lasts: If True, preserves the max_l_size last elements\n",
        "    of a sequence (by keeping the same order).  E.g.:\n",
        "    if keep_lasts is True and max_l_size=3 [1,2,3,4] becomes [2,3,4].\n",
        "\n",
        "\n",
        "    :return: the list padded or truncated.\n",
        "    \"\"\"\n",
        "    to_pad = []\n",
        "    max_l = min(max_l_size, len(l))  # maximum len\n",
        "    l_init = len(l) - max_l if len(l) > max_l and keep_lasts else 0  # initial position where to sample from the list\n",
        "    l_end = len(l) if len(l) > max_l and keep_lasts else max_l\n",
        "    for i in range(l_init, l_end):\n",
        "        to_pad.append(l[i])\n",
        "\n",
        "    # for j in range(len(l), max_l_size):\n",
        "    #     to_pad.append(pad_token)\n",
        "    pad_tokens = [pad_token] * (max_l_size-len(l))\n",
        "    padded_l = to_pad + pad_tokens if pad_right else pad_tokens + to_pad\n",
        "\n",
        "    return padded_l\n",
        "\n",
        "\n",
        "def save_data(data, file):\n",
        "    with open(file, 'wb') as output:\n",
        "        pickle.dump(data, output, pickle.HIGHEST_PROTOCOL)\n",
        "\n",
        "def load_data(file):\n",
        "    with open(file, 'rb') as obj:\n",
        "        return pickle.load(obj)\n",
        "\n",
        "def print_and_write(file, s):\n",
        "    print(s)\n",
        "    file.write(s)\n",
        "\n",
        "\n",
        "class Vocabulary(object):\n",
        "    def __init__(self, vocab_size=None):\n",
        "        self.dictionary = dict()\n",
        "        self.rev_dictionary = dict()\n",
        "        self.count = []\n",
        "        self.special_tokens = []\n",
        "        self.vocab_size = vocab_size\n",
        "\n",
        "    def build_vocabulary_from_counts(self, count, special_tokens=[]):\n",
        "        \"\"\"\n",
        "        Sets all the attributes of the Vocabulary object.\n",
        "        :param count: a list of lists as follows: [['token', number_of_occurrences],...]\n",
        "        :param special_tokens: a list of strings. E.g. ['<EOS>', '<PAD>',...]\n",
        "        :return: None\n",
        "        \"\"\"\n",
        "\n",
        "        dictionary = dict()\n",
        "        for word, _ in count:\n",
        "            dictionary[word] = len(dictionary)\n",
        "\n",
        "        # adding eventual special tokens to the dictionary (e.g. <EOS>,<PAD> etc..)\n",
        "        d = len(dictionary)\n",
        "        for i, token in enumerate(special_tokens):\n",
        "            dictionary[token] = d + i\n",
        "\n",
        "        self.count = count\n",
        "        self.dictionary = dictionary\n",
        "        self.rev_dictionary = dict(zip(self.dictionary.values(), self.dictionary.keys()))\n",
        "        self.special_tokens = special_tokens\n",
        "        self.vocab_size = len(dictionary)\n",
        "\n",
        "    def build_vocabulary_from_tokens(self, tokens, vocabulary_size=None, special_tokens=[]):\n",
        "        \"\"\"\n",
        "        Given a list of tokens, it sets the Vocabulary object attributes by constructing\n",
        "        a dictionary mapping each token to a unique id.\n",
        "        :param tokens: a list of strings.\n",
        "         E.g. [\"the\", \"cat\", \"is\", ... \".\", \"the\", \"house\" ,\"is\" ...].\n",
        "         NB: Here you should put all your token instances of the corpus.\n",
        "        :param vocabulary_size: The number of elements of your vocabulary. If there are more\n",
        "        than 'vocabulary_size' elements on tokens, it considers only the 'vocabulary_size'\n",
        "        most frequent ones.\n",
        "        :param special_tokens: Optional. A list of strings. Useful to add special tokens in vocabulary.\n",
        "        If you don't have any, keep it empty.\n",
        "        :return: None\n",
        "        \"\"\"\n",
        "\n",
        "        vocabulary_size = vocabulary_size if vocabulary_size is not None else self.vocab_size\n",
        "        vocabulary_size = vocabulary_size - (len(special_tokens) + 1) if vocabulary_size else None\n",
        "        # counts occurrences of each token\n",
        "        count = [['<UNK>', -1]]\n",
        "        count.extend(collections.Counter(tokens).most_common(vocabulary_size))  # takes only the most frequent ones, if size is None takes them all\n",
        "        self.build_vocabulary_from_counts(count, special_tokens)  # actually build the vocabulary\n",
        "        self._set_unk_count(tokens)  # set the number of OOV instances\n",
        "\n",
        "    @staticmethod\n",
        "    def merge_vocabulary(vocab0, vocab1, vocabulary_size=-1):\n",
        "        \"\"\"\n",
        "        Merge two Vocabulary objects into a new one.\n",
        "        :param vocab0: first Vocabulary object\n",
        "        :param vocab1: second Vocabulary object\n",
        "        :param vocabulary_size: parameter to decide the merged vocabulary size.\n",
        "        With default value -1, all the words of both vocabularies are preserved.\n",
        "        When set to 0, the size of the vocabulary is set to the size of vocab0,\n",
        "        when set to 1 it is kept the size of vocab1.\n",
        "        :return: a new vocabulary\n",
        "        \"\"\"\n",
        "        # get size of the new vocabulary\n",
        "        vocab_size = vocab0.vocab_size + vocab1.vocab_size if vocabulary_size == -1 else vocabulary_size\n",
        "        merged_special_tokens = list(set(vocab0.special_tokens) | set(vocab1.special_tokens))\n",
        "\n",
        "        # merge the counts from the two vocabularies and then selects the most_common tokens\n",
        "        merged_counts = collections.Counter(dict(vocab0.count)) + collections.Counter(dict(vocab1.count))\n",
        "        merged_counts = merged_counts.most_common(vocab_size)\n",
        "        count = [['<UNK>', -1]]\n",
        "        count.extend(merged_counts)\n",
        "\n",
        "        # create the new vocabulary\n",
        "        merged_vocab = Vocabulary(vocab_size)\n",
        "        merged_vocab.build_vocabulary_from_counts(count, merged_special_tokens)\n",
        "        return merged_vocab\n",
        "\n",
        "    @staticmethod\n",
        "    def merge_vocabularies(vocab_list, vocab_size=None):\n",
        "        \"\"\"\n",
        "        Join a list of vocabularies into a new one.\n",
        "        :param vocab_list: a list of Vocabulary objects\n",
        "        :param vocab_size: the maximum size of the merged vocabulary.\n",
        "        :return: a vocabulary merging them all.\n",
        "        \"\"\"\n",
        "        vocab_size = vocab_size if vocab_size else sum([v.vocab_size for v in vocab_list])\n",
        "        merged_vocab = Vocabulary(vocab_size)\n",
        "        for voc in vocab_list:\n",
        "            merged_vocab = Vocabulary.merge_vocabulary(merged_vocab, voc, vocab_size)\n",
        "        return merged_vocab\n",
        "\n",
        "    def string2id(self, dataset):\n",
        "        \"\"\"\n",
        "        Converts a dataset of strings into a dataset of ids according to the object dictionary.\n",
        "        :param dataset: any string-based dataset with any nested lists.\n",
        "        :return: a new dataset, with the same shape of dataset, where each string is mapped into its\n",
        "        corresponding id associated in the dictionary (0 for unknown tokens).\n",
        "        \"\"\"\n",
        "\n",
        "        def _recursive_call(items):\n",
        "            new_items = []\n",
        "            for item in items:\n",
        "                if isinstance(item, str) or isinstance(item, int) or isinstance(item, float):\n",
        "                    new_items.append(self.word2id(item))\n",
        "                else:\n",
        "                    new_items.append(_recursive_call(item))\n",
        "            return new_items\n",
        "\n",
        "        return _recursive_call(dataset)\n",
        "\n",
        "    def id2string(self, dataset):\n",
        "        \"\"\"\n",
        "        Converts a dataset of integer ids into a dataset of string according to the reverse dictionary.\n",
        "        :param dataset: any int-based dataset with any nested lists. Allowed types are int, np.int32, np.int64.\n",
        "        :return: a new dataset, with the same shape of dataset, where each token is mapped into its\n",
        "        corresponding string associated in the reverse dictionary.\n",
        "        \"\"\"\n",
        "        def _recursive_call(items):\n",
        "            new_items = []\n",
        "            for item in items:\n",
        "                if isinstance(item, int) or isinstance(item, np.int) or isinstance(item, np.int32) or isinstance(item, np.int64):\n",
        "                    new_items.append(self.id2word(item))\n",
        "                else:\n",
        "                    new_items.append(_recursive_call(item))\n",
        "            return new_items\n",
        "\n",
        "        return _recursive_call(dataset)\n",
        "\n",
        "    def word2id(self, item):\n",
        "        \"\"\"\n",
        "        Maps a string token to its corresponding id.\n",
        "        :param item: a string.\n",
        "        :return: If the token belongs to the vocabulary, it returns an integer id > 0, otherwise\n",
        "        it returns the value associated to the unknown symbol, that is typically 0.\n",
        "        \"\"\"\n",
        "        return self.dictionary[item] if item in self.dictionary else self.dictionary['<UNK>']\n",
        "\n",
        "    def id2word(self, token_id):\n",
        "        \"\"\"\n",
        "        Maps an integer token to its corresponding string.\n",
        "        :param token_id: an integer.\n",
        "        :return: If the id belongs to the vocabulary, it returns the string\n",
        "        associated to it, otherwise it returns the string associated\n",
        "        to the unknown symbol, that is '<UNK>'.\n",
        "        \"\"\"\n",
        "\n",
        "        return self.rev_dictionary[token_id] if token_id in self.rev_dictionary else self.rev_dictionary[self.dictionary['<UNK>']]\n",
        "\n",
        "    def get_unk_count(self):\n",
        "        return self.count[0][1]\n",
        "\n",
        "    def _set_unk_count(self, tokens):\n",
        "        \"\"\"\n",
        "        Sets the number of OOV instances in the tokens provided\n",
        "        :param tokens: a list of tokens\n",
        "        :return: None\n",
        "        \"\"\"\n",
        "        data = list()\n",
        "        unk_count = 0\n",
        "        for word in tokens:\n",
        "            if word in self.dictionary:\n",
        "                index = self.dictionary[word]\n",
        "            else:\n",
        "                index = 0  # dictionary['<UNK>']\n",
        "                unk_count += 1\n",
        "            data.append(index)\n",
        "        self.count[0][1] = unk_count\n",
        "\n",
        "    def add_element(self, name, is_special_token=False):\n",
        "        if name not in self.dictionary:\n",
        "            self.vocab_size += 1\n",
        "            self.dictionary[name] = self.vocab_size\n",
        "            self.rev_dictionary[self.vocab_size] = name\n",
        "\n",
        "            if is_special_token:\n",
        "                self.special_tokens = list(self.special_tokens)\n",
        "                self.special_tokens.append(name)\n",
        "\n",
        "            self.count.append([name, 1])\n",
        "\n",
        "    def set_vocabulary(self, dictionary, rev_dictionary, special_tokens, vocab_size):\n",
        "        self.dictionary = dictionary,\n",
        "        self.rev_dictionary = rev_dictionary\n",
        "        self.special_tokens = special_tokens\n",
        "        self.vocab_size = vocab_size\n",
        "\n",
        "    @staticmethod\n",
        "    def load_vocabulary(filename):\n",
        "        return load_data(filename)\n",
        "\n",
        "    def save_vocabulary(self, filename):\n",
        "        save_data(self, filename)\n",
        "\n",
        "class SyLMDataset(object):\n",
        "    def __init__(self, config, sy_vocab=None):\n",
        "        self.config = config\n",
        "        self.vocabulary = sy_vocab\n",
        "\n",
        "        self.raw_train_x = []\n",
        "        self.raw_val_x = []\n",
        "        self.raw_test_x = []\n",
        "        self.raw_x = []\n",
        "\n",
        "        self.train_x, self.train_y = [], []\n",
        "        self.val_x, self.val_y = [], []\n",
        "        self.test_x, self.test_y = [], []\n",
        "        self.x, self.y = [], []\n",
        "\n",
        "    def initialize(self, sess):\n",
        "        pass\n",
        "\n",
        "    def load(self, sources):\n",
        "        \"\"\"\n",
        "        Extract raw texts form sources and gather them all together.\n",
        "        :param sources: a string or an iterable of strings containing the file(s)\n",
        "        to process in order to build the dataset.\n",
        "        :return: a list of raw strings.\n",
        "        \"\"\"\n",
        "        return NotImplementedError\n",
        "\n",
        "    def build(self, sources, split_size=0.8):\n",
        "        \"\"\"\n",
        "        :param sources: a string or an iterable of strings containing the file(s)\n",
        "        to process in order to build the dataset.\n",
        "        :param split_size: the size to split the dataset, set >=1.0 to not split.\n",
        "        \"\"\"\n",
        "\n",
        "        raw_x = self.load(sources)\n",
        "        # raw_x = self.tokenize([self.preprocess(ex) for ex in raw_x])  # fixme\n",
        "        # splitting data\n",
        "        self.raw_x = raw_x\n",
        "        if split_size < 1.0:\n",
        "            self.raw_train_x, self.raw_test_x = self.split(self.raw_x, train_size=split_size)\n",
        "            self.raw_train_x, self.raw_val_x = self.split(self.raw_train_x, train_size=split_size)\n",
        "        else:\n",
        "            self.raw_train_x = self.raw_x\n",
        "\n",
        "        if self.vocabulary is None:\n",
        "            # creates vocabulary\n",
        "            tokens = [item for sublist in self.raw_train_x for item in sublist]  # get tokens\n",
        "            special_tokens = (\"<GO>\", \"<PAD>\", \"<SEP>\", \"<EOS>\", \"<EOV>\")\n",
        "            self._create_vocab(tokens, special_tokens=special_tokens)\n",
        "\n",
        "        # creates x,y for train\n",
        "        self.train_x = self._build_dataset(self.raw_train_x, insert_go=True, max_len=self.config.sentence_max_len, shuffle=False)\n",
        "        self.train_y = self._build_dataset(self.raw_train_x, insert_go=True, max_len=self.config.sentence_max_len, shuffle=False)\n",
        "\n",
        "        # creates x,y for validation\n",
        "        self.val_x = self._build_dataset(self.raw_val_x, insert_go=True, max_len=self.config.sentence_max_len, shuffle=False)\n",
        "        self.val_y = self._build_dataset(self.raw_val_x, insert_go=True, max_len=self.config.sentence_max_len, shuffle=False)\n",
        "\n",
        "        # creates x,y for validation\n",
        "        self.test_x = self._build_dataset(self.raw_test_x, insert_go=True, max_len=self.config.sentence_max_len, shuffle=False)\n",
        "        self.test_y = self._build_dataset(self.raw_test_x, insert_go=True, max_len=self.config.sentence_max_len, shuffle=False)\n",
        "\n",
        "    def _create_vocab(self, tokens, special_tokens=(\"<PAD>\", \"<GO>\", \"<SEP>\", \"<EOV>\", \"<EOS>\")):\n",
        "        \"\"\"\n",
        "        Create the vocabulary. Special tokens can be added to the tokens obtained from\n",
        "        the corpus.\n",
        "        :param tokens: a list of all the tokens in the corpus. Each token is a string.\n",
        "        :param special_tokens: a list of strings.\n",
        "        \"\"\"\n",
        "        vocab = Vocabulary(vocab_size=self.config.input_vocab_size)\n",
        "        vocab.build_vocabulary_from_tokens(tokens, special_tokens=special_tokens)\n",
        "        self.vocabulary = vocab\n",
        "\n",
        "    @staticmethod\n",
        "    def split(raw_data, train_size=0.8):\n",
        "        size = math.floor(len(raw_data)*train_size)\n",
        "        return raw_data[:size], raw_data[size:]\n",
        "\n",
        "    @staticmethod\n",
        "    def preprocess(txt):\n",
        "        return txt\n",
        "\n",
        "    @staticmethod\n",
        "    def shuffle(x):\n",
        "        return random.sample(x, len(x))\n",
        "\n",
        "    @staticmethod\n",
        "    def tokenize(txt):\n",
        "        return txt\n",
        "\n",
        "    def _build_dataset(self, raw_data, max_len=100, insert_go=True, keep_lasts=False, pad_right=True, shuffle=True):\n",
        "        \"\"\"\n",
        "        Converts all the tokens in e1_raw_data by mapping each token with its corresponding\n",
        "        value in the dictionary. In case of token not in the dictionary, they are assigned to\n",
        "        a specific id. Each sequence is padded up to the seq_max_len setup in the config.\n",
        "\n",
        "        :param raw_data: list of sequences, each sequence is a list of tokens (strings).\n",
        "        :param max_len: max length of a sequence, crop longer and pad smaller ones.\n",
        "        :param insert_go: True to insert <GO>, False otherwise.\n",
        "        :param keep_lasts: True to truncate initial elements of a sequence.\n",
        "        :param pad_right: pad to the right (default value True), otherwise pads to left.\n",
        "        :param shuffle: Optional. If True data are shuffled.\n",
        "        :return: A list of sequences where each token in each sequence is an int id.\n",
        "        \"\"\"\n",
        "        dataset = []\n",
        "        for sentence in raw_data:\n",
        "            sentence_ids = [self.vocabulary.word2id(\"<GO>\")] if insert_go else []\n",
        "            sentence_ids.extend([self.vocabulary.word2id(w) for w in sentence])\n",
        "            sentence_ids.append(self.vocabulary.word2id(\"<EOS>\"))\n",
        "            sentence_ids = pad_list(sentence_ids, self.vocabulary.word2id(\"<PAD>\"), max_len, keep_lasts=keep_lasts, pad_right=pad_right)\n",
        "\n",
        "            dataset.append(sentence_ids)\n",
        "\n",
        "        if shuffle:\n",
        "            return random.sample(dataset, len(dataset))\n",
        "        else:\n",
        "            return dataset\n",
        "\n",
        "    def get_batches(self, batch_size=32, split_sel='train'):\n",
        "        \"\"\"\n",
        "        Iterator over the training set. Useful method to run experiments.\n",
        "        :param batch_size: size of the mini_batch\n",
        "        :return: input and target.\n",
        "        \"\"\"\n",
        "        if split_sel == 'train':\n",
        "            x, y = self.train_x, self.train_y\n",
        "        elif split_sel == 'val':\n",
        "            x, y = self.val_x, self.val_y\n",
        "        else:\n",
        "            x, y = self.test_x, self.test_y\n",
        "        \n",
        "        i = 0#random.randint(0, batch_size)\n",
        "        batches = []\n",
        "        eov = self.vocabulary.word2id(\"<EOV>\")\n",
        "        go = self.vocabulary.word2id(\"<GO>\")\n",
        "        # prepare batches\n",
        "        while i < len(x):\n",
        "            j = 0\n",
        "            batch_x, batch_y = [], []\n",
        "            while j < batch_size and i+j<len(x):\n",
        "                for c in x[i+j]:\n",
        "                  batch_x.append(c)\n",
        "                batch_x.append(eov)\n",
        "                for c in y[i+j]:\n",
        "                  batch_y.append(c)\n",
        "                batch_y.append(eov)\n",
        "                j += 1\n",
        "            i += batch_size\n",
        "            batches.append((batch_x, batch_y))\n",
        "\n",
        "        # supply\n",
        "        i = 0\n",
        "        while i < len(batches):\n",
        "            yield batches[i][0], batches[i][1]\n",
        "            i += 1\n",
        "\n",
        "class DanteSyLMDataset(SyLMDataset):\n",
        "    def __init__(self, config, sy_vocab=None):\n",
        "        \"\"\"\n",
        "        Class to create a dataset from Dante Alighieri's Divine Comedy.\n",
        "        :param config: a Config object\n",
        "        :param sy_vocab: (optional) a Vocabulary object where tokens of the dictionary\n",
        "        are syllables. If None, the vocabulary is create automatically from the source.\n",
        "        \"\"\"\n",
        "        super().__init__(config, sy_vocab)\n",
        "\n",
        "    def load(self, sources):\n",
        "        \"\"\"\n",
        "        Load examples from dataset\n",
        "        :param sources: data filepath.\n",
        "        :return:\n",
        "        \"\"\"\n",
        "        canti, _, raw = get_dc_cantos(filename=sources)  # get raw data from file\n",
        "        canti, tokens = get_dc_hyphenation(canti)  # converts each\n",
        "\n",
        "        tercets = create_tercets(canti)\n",
        "        tercets = get_hyp_lm_tercets(tercets)\n",
        "\n",
        "        x = []\n",
        "        for tercet in tercets:\n",
        "            x.append([])\n",
        "            for verse in tercet:\n",
        "                x[-1].extend(verse)\n",
        "                x[-1].append(\"<EOV>\")\n",
        "\n",
        "        #x = self.shuffle(x)\n",
        "        return x\n",
        "\n",
        "def seq2str(seq):\n",
        "    def output2string(batch, rev_vocabulary, special_tokens, end_of_tokens):\n",
        "        to_print = ''\n",
        "        for token in batch:\n",
        "            if token in special_tokens:\n",
        "                to_print += ' '\n",
        "            elif end_of_tokens and token in end_of_tokens:\n",
        "                to_print += '\\n'\n",
        "            elif token in rev_vocabulary:\n",
        "                to_print += rev_vocabulary[token]\n",
        "            else:\n",
        "                to_print += '<UNK>'\n",
        "        return to_print\n",
        "\n",
        "    return output2string(seq, poetry_sy_lm_dataset.vocabulary.rev_dictionary,\n",
        "      special_tokens=[poetry_sy_lm_dataset.vocabulary.word2id(\"<PAD>\"), 0, poetry_sy_lm_dataset.vocabulary.word2id(\"<SEP>\"),\n",
        "                      poetry_sy_lm_dataset.vocabulary.word2id(\"<GO>\"), poetry_sy_lm_dataset.vocabulary.word2id(\"<EOS>\")],\n",
        "      end_of_tokens=[poetry_sy_lm_dataset.vocabulary.word2id(\"<EOV>\")])\n",
        "\n",
        "class cnfg:\n",
        "  vocab_size = vocab_size\n",
        "  input_vocab_size = vocab_size\n",
        "  sentence_max_len = terces_len\n",
        "\n",
        "config = cnfg()\n",
        "poetry_sy_lm_dataset = DanteSyLMDataset(config, sy_vocab=None)\n",
        "url = \"https://raw.githubusercontent.com/DanieleVeri/deep_comedy/master/divina_commedia.txt\"\n",
        "response = requests.get(url)\n",
        "#response.encoding = 'ISO-8859-1'\n",
        "fi = open(\"divcom.txt\",\"w\")\n",
        "fi.write(response.text)\n",
        "fi.close()\n",
        "data_path = os.path.join(os.getcwd(), \"divcom.txt\")  # dataset location, here just the name of the source file\n",
        "poetry_sy_lm_dataset.build(data_path, split_size=0.99)  # actual creation of  vocabulary (if not provided) and dataset\n",
        "print(\"Train size: \" + str(len(poetry_sy_lm_dataset.train_y)))\n",
        "print(\"Val size: \" + str(len(poetry_sy_lm_dataset.val_y)))\n",
        "print(\"Test size: \" + str(len(poetry_sy_lm_dataset.test_y)))\n",
        "\n",
        "eov = poetry_sy_lm_dataset.vocabulary.word2id(\"<EOV>\")\n",
        "pad = poetry_sy_lm_dataset.vocabulary.word2id(\"<PAD>\")\n",
        "go = poetry_sy_lm_dataset.vocabulary.word2id(\"<GO>\")\n",
        "eos = poetry_sy_lm_dataset.vocabulary.word2id(\"<EOS>\")\n",
        "\n",
        "batches = [b for b in poetry_sy_lm_dataset.get_batches(terces_per_batch)]\n",
        "print(batches[0][0])\n",
        "print(batches[0][1])\n",
        "print(len(batches[0][0]))\n",
        "test_b = [b for b in poetry_sy_lm_dataset.get_batches(terces_per_batch, split_sel='test')]\n",
        "print(test_b[0][0])\n",
        "print(test_b[0][1])\n",
        "print(len(test_b[0][0]))\n",
        "val_b = [b for b in poetry_sy_lm_dataset.get_batches(terces_per_batch, split_sel='val')]\n",
        "print(val_b[0][0])\n",
        "print(val_b[0][1])\n",
        "print(len(val_b[0][0]))\n",
        "len(poetry_sy_lm_dataset.vocabulary.dictionary.items())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sEun8YPR8Sn6",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Generation\n",
        "def generate(index=0):\n",
        "\n",
        "    def evaluate_greedy(inp_sentence, decoder_input):\n",
        "        inp_sentence = inp_sentence\n",
        "        encoder_input = tf.expand_dims(inp_sentence, 0)\n",
        "        \n",
        "        output = tf.expand_dims(decoder_input, 0)\n",
        "\n",
        "        terces = 0\n",
        "        for i in range(batch_len):\n",
        "            enc_padding_mask, combined_mask, dec_padding_mask = create_masks(\n",
        "                encoder_input, output)\n",
        "        \n",
        "            # predictions.shape == (batch_size, seq_len, vocab_size)\n",
        "            predictions, attention_weights = transformer(encoder_input, \n",
        "                                                        output,\n",
        "                                                        False,\n",
        "                                                        enc_padding_mask,\n",
        "                                                        combined_mask,\n",
        "                                                        dec_padding_mask)\n",
        "            \n",
        "            # select the last word from the seq_len dimension\n",
        "            predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)\n",
        "\n",
        "            predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)\n",
        "\n",
        "            # return the result if the predicted_id is equal to the end token\n",
        "            if predicted_id == eos:\n",
        "                terces += 1\n",
        "                if terces == terces_per_batch-1:\n",
        "                    return tf.squeeze(output, axis=0), attention_weights\n",
        "            # concatentate the predicted_id to the output which is given to the decoder\n",
        "            # as its input.\n",
        "            output = tf.concat([output, predicted_id], axis=-1)\n",
        "\n",
        "        return tf.squeeze(output, axis=0), attention_weights\n",
        "\n",
        "\n",
        "    def evaluate_topk(inp_sentence, decoder_input, k=5, temperature=0.5):\n",
        "        inp_sentence = inp_sentence\n",
        "        encoder_input = tf.expand_dims(inp_sentence, 0)\n",
        "        \n",
        "        output = tf.expand_dims(decoder_input, 0)\n",
        "\n",
        "        def scale(tensor):\n",
        "            tensor = tf.math.divide(\n",
        "                tf.subtract(\n",
        "                    tensor, \n",
        "                    tf.reduce_min(tensor)\n",
        "                ), \n",
        "                tf.subtract(\n",
        "                    tf.reduce_max(tensor), \n",
        "                    tf.reduce_min(tensor))\n",
        "                )\n",
        "            return tensor\n",
        "\n",
        "        terces = 0\n",
        "        for i in range(batch_len):\n",
        "            enc_padding_mask, combined_mask, dec_padding_mask = create_masks(\n",
        "                encoder_input, output)\n",
        "        \n",
        "            # predictions.shape == (batch_size, seq_len, vocab_size)\n",
        "            predictions, attention_weights = transformer(encoder_input, \n",
        "                                                        output,\n",
        "                                                        False,\n",
        "                                                        enc_padding_mask,\n",
        "                                                        combined_mask,\n",
        "                                                        dec_padding_mask)\n",
        "            # select the last word from the seq_len dimension\n",
        "            predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)\n",
        "            predictions, indices = tf.math.top_k(predictions,k=k)\n",
        "            predictions /= temperature\n",
        "            #predictions = scale(predictions)\n",
        "            predictions = np.squeeze(predictions, axis=0)\n",
        "            indices = np.squeeze(indices, axis=0)\n",
        "            indices = np.squeeze(indices, axis=0)\n",
        "            predicted_id = tf.random.categorical(predictions, num_samples=1)[-1,0].numpy()\n",
        "            predicted_id = indices[predicted_id]\n",
        "\n",
        "            # return the result if the predicted_id is equal to the end token\n",
        "            if predicted_id == eos:\n",
        "                terces += 1\n",
        "                if terces == terces_per_batch-1:\n",
        "                    return tf.squeeze(output, axis=0), attention_weights\n",
        "            # concatentate the predicted_id to the output which is given to the decoder\n",
        "            # as its input.\n",
        "            predicted_id = tf.expand_dims(predicted_id, 0)\n",
        "            predicted_id = tf.expand_dims(predicted_id, 0)\n",
        "            output = tf.concat([output, predicted_id], axis=-1)\n",
        "\n",
        "        return tf.squeeze(output, axis=0), attention_weights\n",
        "\n",
        "    out_list = test_b[index][0]\n",
        "    offset = terces_len # a tercet\n",
        "    txt_gen = seq2str(out_list[-offset:])\n",
        "    k = 1\n",
        "    t = 1\n",
        "    for i in range(32//(terces_per_batch-1)): # 30 terces = cantica\n",
        "        out, att_w = evaluate_topk([pad], out_list[-offset:], k, t)\n",
        "        out_list = out.numpy().tolist()\n",
        "        out_str = seq2str(out_list[offset:])\n",
        "        txt_gen += out_str\n",
        "\n",
        "    print(txt_gen)\n",
        "    wandb.log({\"generated\":\n",
        "            wandb.Html(\"k=\"+str(k)+\" t=\"+str(t)+\n",
        "                       \"<pre>\"+txt_gen+\"</pre>\", inject=False)})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZwdCuj3HsFux",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        " #@title Train loop\n",
        "\n",
        "for epoch in range(EPOCHS):\n",
        "    random.shuffle(batches)\n",
        "    start = time.time()\n",
        "    \n",
        "    train_loss.reset_states()\n",
        "    train_accuracy.reset_states()\n",
        "    \n",
        "    for (batch, (inp, tar)) in enumerate(batches):\n",
        "        if (len(inp) != batch_len or len(tar) != batch_len):\n",
        "            print(\"discarded batch\", batch)\n",
        "            continue\n",
        "        train_step(np.expand_dims(inp, axis=0), np.expand_dims(tar, axis=0))\n",
        "        \n",
        "        if batch % 50 == 0:\n",
        "            print ('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(\n",
        "                epoch + 1, batch, train_loss.result(), train_accuracy.result()))\n",
        "        \n",
        "    if (epoch + 1) % 5 == 0:\n",
        "        ckpt_save_path = ckpt_manager.save()\n",
        "        print ('Saving checkpoint for epoch {} at {}'.format(epoch+1, ckpt_save_path))\n",
        "\n",
        "    print ('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1, train_loss.result(), train_accuracy.result()))\n",
        "    print ('Time taken for 1 epoch: {} secs\\n'.format(time.time() - start))\n",
        "\n",
        "    wandb.log({\n",
        "        'train_loss': train_loss.result(),\n",
        "        'train_accuracy': train_accuracy.result()\n",
        "    }, step=epoch+1)\n",
        "\n",
        "    # validation\n",
        "    if epoch % 5 == 0:\n",
        "        loss_l, acc_l = [], []\n",
        "        for (batch, (inp, tar)) in enumerate(val_b):\n",
        "            val_loss.reset_states()\n",
        "            val_accuracy.reset_states()\n",
        "            \n",
        "            if (len(inp) != batch_len or len(tar) != batch_len):\n",
        "                print(\"discarded batch\", batch)\n",
        "                continue\n",
        "\n",
        "            val_step(np.expand_dims(inp, axis=0), np.expand_dims(tar, axis=0))\n",
        "\n",
        "            loss_l.append(val_loss.result())\n",
        "            acc_l.append(val_accuracy.result())\n",
        "\n",
        "        loss_mean = sum(loss_l)/len(loss_l)\n",
        "        acc_mean = sum(acc_l)/len(acc_l)\n",
        "        print('Epoch {} VALIDATION: Loss {:.4f} Accuracy {:.4f}\\n'.format(epoch + 1, loss_mean, acc_mean))\n",
        "\n",
        "        wandb.log({\n",
        "            'val_loss': loss_mean,\n",
        "            'val_accuracy': acc_mean\n",
        "        }, step=epoch+1)\n",
        "\n",
        "    # generation\n",
        "    if epoch in generate_at:\n",
        "        generate(1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SjLbfLOU_pRw",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Parameter persistence\n",
        "\n",
        "transformer.save_weights(\"./optimus_rhyme\")\n",
        "#transformer.load_weights(\"./optimus_rhyme\")\n",
        "\n",
        "\n",
        "#emb_enc_w = transformer.encoder.embedding.get_weights()[0]\n",
        "emb_enc_w = transformer.decoder.embedding.get_weights()[0]\n",
        "print(emb_enc_w.shape)\n",
        "\n",
        "out_v = open('vecs.tsv', 'w', encoding='utf-8')\n",
        "out_m = open('meta.tsv', 'w', encoding='utf-8')\n",
        "\n",
        "for num, word in enumerate(poetry_sy_lm_dataset.vocabulary.dictionary):\n",
        "  vec = emb_enc_w[num] # skip 0, it's padding.\n",
        "  out_m.write(word + \"\\n\")\n",
        "  out_v.write('\\t'.join([str(x) for x in vec]) + \"\\n\")\n",
        "out_v.close()\n",
        "out_m.close()\n",
        "\n",
        "\n",
        "'''\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir .\n",
        "'''"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}