NLP/Text_Generation_using_FNet.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Text_Generation_using_FNet.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "RwkH0fHnIOj0"
      },
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "import os\n",
        "import re\n",
        "\n",
        "# Defining hyperparameters\n",
        "\n",
        "VOCAB_SIZE = 8192\n",
        "MAX_SAMPLES = 50000\n",
        "BUFFER_SIZE = 20000\n",
        "MAX_LENGTH = 40\n",
        "EMBED_DIM = 256\n",
        "LATENT_DIM = 512\n",
        "NUM_HEADS = 8\n",
        "BATCH_SIZE = 64"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "I1wtWVkyIsjG",
        "outputId": "a745d8d3-401d-4a7a-8efd-691f5def5704"
      },
      "source": [
        "path_to_zip = keras.utils.get_file(\n",
        "    \"cornell_movie_dialogs.zip\",\n",
        "    origin=\"http://www.cs.cornell.edu/~cristian/data/cornell_movie_dialogs_corpus.zip\",\n",
        "    extract=True,\n",
        ")\n",
        "\n",
        "path_to_dataset = os.path.join(\n",
        "    os.path.dirname(path_to_zip), \"cornell movie-dialogs corpus\"\n",
        ")\n",
        "path_to_movie_lines = os.path.join(path_to_dataset, \"movie_lines.txt\")\n",
        "path_to_movie_conversations = os.path.join(path_to_dataset, \"movie_conversations.txt\")\n",
        "\n",
        "\n",
        "def load_conversations():\n",
        "    # Helper function for loading the conversation splits\n",
        "    id2line = {}\n",
        "    with open(path_to_movie_lines, errors=\"ignore\") as file:\n",
        "        lines = file.readlines()\n",
        "    for line in lines:\n",
        "        parts = line.replace(\"\\n\", \"\").split(\" +++$+++ \")\n",
        "        id2line[parts[0]] = parts[4]\n",
        "\n",
        "    inputs, outputs = [], []\n",
        "    with open(path_to_movie_conversations, \"r\") as file:\n",
        "        lines = file.readlines()\n",
        "    for line in lines:\n",
        "        parts = line.replace(\"\\n\", \"\").split(\" +++$+++ \")\n",
        "        # get conversation in a list of line ID\n",
        "        conversation = [line[1:-1] for line in parts[3][1:-1].split(\", \")]\n",
        "        for i in range(len(conversation) - 1):\n",
        "            inputs.append(id2line[conversation[i]])\n",
        "            outputs.append(id2line[conversation[i + 1]])\n",
        "            if len(inputs) >= MAX_SAMPLES:\n",
        "                return inputs, outputs\n",
        "    return inputs, outputs\n",
        "\n",
        "\n",
        "questions, answers = load_conversations()\n",
        "\n",
        "# Splitting training and validation sets\n",
        "\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices((questions[:40000], answers[:40000]))\n",
        "val_dataset = tf.data.Dataset.from_tensor_slices((questions[40000:], answers[40000:]))"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading data from http://www.cs.cornell.edu/~cristian/data/cornell_movie_dialogs_corpus.zip\n",
            "9920512/9916637 [==============================] - 2s 0us/step\n",
            "9928704/9916637 [==============================] - 2s 0us/step\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q1h-_VEvIxKP"
      },
      "source": [
        "def preprocess_text(sentence):\n",
        "    sentence = tf.strings.lower(sentence)\n",
        "    # Adding a space between the punctuation and the last word to allow better tokenization\n",
        "    sentence = tf.strings.regex_replace(sentence, r\"([?.!,])\", r\" \\1 \")\n",
        "    # Replacing multiple continuous spaces with a single space\n",
        "    sentence = tf.strings.regex_replace(sentence, r\"\\s\\s+\", \" \")\n",
        "    # Replacing non english words with spaces\n",
        "    sentence = tf.strings.regex_replace(sentence, r\"[^a-z?.!,]+\", \" \")\n",
        "    sentence = tf.strings.strip(sentence)\n",
        "    sentence = tf.strings.join([\"[start]\", sentence, \"[end]\"], separator=\" \")\n",
        "    return sentence\n",
        "\n",
        "\n",
        "vectorizer = layers.TextVectorization(\n",
        "    VOCAB_SIZE,\n",
        "    standardize=preprocess_text,\n",
        "    output_mode=\"int\",\n",
        "    output_sequence_length=MAX_LENGTH,\n",
        ")\n",
        "\n",
        "# We will adapt the vectorizer to both the questions and answers\n",
        "# This dataset is batched to parallelize and speed up the process\n",
        "vectorizer.adapt(tf.data.Dataset.from_tensor_slices((questions + answers)).batch(128))"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "of8I2vbeI0ms"
      },
      "source": [
        "def vectorize_text(inputs, outputs):\n",
        "    inputs, outputs = vectorizer(inputs), vectorizer(outputs)\n",
        "    # One extra padding token to the right to match the output shape\n",
        "    outputs = tf.pad(inputs, [[0, 1]])\n",
        "    return (\n",
        "        {\"encoder_inputs\": inputs, \"decoder_inputs\": outputs[:-1]},\n",
        "        {\"outputs\": outputs[1:]},\n",
        "    )\n",
        "\n",
        "\n",
        "train_dataset = train_dataset.map(vectorize_text, num_parallel_calls=tf.data.AUTOTUNE)\n",
        "val_dataset = val_dataset.map(vectorize_text, num_parallel_calls=tf.data.AUTOTUNE)\n",
        "\n",
        "train_dataset = (\n",
        "    train_dataset.cache()\n",
        "    .shuffle(BUFFER_SIZE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.AUTOTUNE)\n",
        ")\n",
        "val_dataset = val_dataset.cache().batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8yU-7FhtI4O9"
      },
      "source": [
        "class FNetEncoder(layers.Layer):\n",
        "    def __init__(self, embed_dim, dense_dim, **kwargs):\n",
        "        super(FNetEncoder, self).__init__(**kwargs)\n",
        "        self.embed_dim = embed_dim\n",
        "        self.dense_dim = dense_dim\n",
        "        self.dense_proj = keras.Sequential(\n",
        "            [\n",
        "                layers.Dense(dense_dim, activation=\"relu\"),\n",
        "                layers.Dense(embed_dim),\n",
        "            ]\n",
        "        )\n",
        "        self.layernorm_1 = layers.LayerNormalization()\n",
        "        self.layernorm_2 = layers.LayerNormalization()\n",
        "\n",
        "    def call(self, inputs):\n",
        "        # Casting the inputs to complex64\n",
        "        inp_complex = tf.cast(inputs, tf.complex64)\n",
        "        # Projecting the inputs to the frequency domain using FFT2D and\n",
        "        # extracting the real part of the output\n",
        "        fft = tf.math.real(tf.signal.fft2d(inp_complex))\n",
        "        proj_input = self.layernorm_1(inputs + fft)\n",
        "        proj_output = self.dense_proj(proj_input)\n",
        "        return self.layernorm_2(proj_input + proj_output)"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U7XmIjm1JGsu"
      },
      "source": [
        "class PositionalEmbedding(layers.Layer):\n",
        "    def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs):\n",
        "        super(PositionalEmbedding, self).__init__(**kwargs)\n",
        "        self.token_embeddings = layers.Embedding(\n",
        "            input_dim=vocab_size, output_dim=embed_dim\n",
        "        )\n",
        "        self.position_embeddings = layers.Embedding(\n",
        "            input_dim=sequence_length, output_dim=embed_dim\n",
        "        )\n",
        "        self.sequence_length = sequence_length\n",
        "        self.vocab_size = vocab_size\n",
        "        self.embed_dim = embed_dim\n",
        "\n",
        "    def call(self, inputs):\n",
        "        length = tf.shape(inputs)[-1]\n",
        "        positions = tf.range(start=0, limit=length, delta=1)\n",
        "        embedded_tokens = self.token_embeddings(inputs)\n",
        "        embedded_positions = self.position_embeddings(positions)\n",
        "        return embedded_tokens + embedded_positions\n",
        "\n",
        "    def compute_mask(self, inputs, mask=None):\n",
        "        return tf.math.not_equal(inputs, 0)\n",
        "\n",
        "\n",
        "class FNetDecoder(layers.Layer):\n",
        "    def __init__(self, embed_dim, latent_dim, num_heads, **kwargs):\n",
        "        super(FNetDecoder, self).__init__(**kwargs)\n",
        "        self.embed_dim = embed_dim\n",
        "        self.latent_dim = latent_dim\n",
        "        self.num_heads = num_heads\n",
        "        self.attention_1 = layers.MultiHeadAttention(\n",
        "            num_heads=num_heads, key_dim=embed_dim\n",
        "        )\n",
        "        self.attention_2 = layers.MultiHeadAttention(\n",
        "            num_heads=num_heads, key_dim=embed_dim\n",
        "        )\n",
        "        self.dense_proj = keras.Sequential(\n",
        "            [\n",
        "                layers.Dense(latent_dim, activation=\"relu\"),\n",
        "                layers.Dense(embed_dim),\n",
        "            ]\n",
        "        )\n",
        "        self.layernorm_1 = layers.LayerNormalization()\n",
        "        self.layernorm_2 = layers.LayerNormalization()\n",
        "        self.layernorm_3 = layers.LayerNormalization()\n",
        "        self.supports_masking = True\n",
        "\n",
        "    def call(self, inputs, encoder_outputs, mask=None):\n",
        "        causal_mask = self.get_causal_attention_mask(inputs)\n",
        "        if mask is not None:\n",
        "            padding_mask = tf.cast(mask[:, tf.newaxis, :], dtype=\"int32\")\n",
        "            padding_mask = tf.minimum(padding_mask, causal_mask)\n",
        "\n",
        "        attention_output_1 = self.attention_1(\n",
        "            query=inputs, value=inputs, key=inputs, attention_mask=causal_mask\n",
        "        )\n",
        "        out_1 = self.layernorm_1(inputs + attention_output_1)\n",
        "\n",
        "        attention_output_2 = self.attention_2(\n",
        "            query=out_1,\n",
        "            value=encoder_outputs,\n",
        "            key=encoder_outputs,\n",
        "            attention_mask=padding_mask,\n",
        "        )\n",
        "        out_2 = self.layernorm_2(out_1 + attention_output_2)\n",
        "\n",
        "        proj_output = self.dense_proj(out_2)\n",
        "        return self.layernorm_3(out_2 + proj_output)\n",
        "\n",
        "    def get_causal_attention_mask(self, inputs):\n",
        "        input_shape = tf.shape(inputs)\n",
        "        batch_size, sequence_length = input_shape[0], input_shape[1]\n",
        "        i = tf.range(sequence_length)[:, tf.newaxis]\n",
        "        j = tf.range(sequence_length)\n",
        "        mask = tf.cast(i >= j, dtype=\"int32\")\n",
        "        mask = tf.reshape(mask, (1, input_shape[1], input_shape[1]))\n",
        "        mult = tf.concat(\n",
        "            [tf.expand_dims(batch_size, -1), tf.constant([1, 1], dtype=tf.int32)],\n",
        "            axis=0,\n",
        "        )\n",
        "        return tf.tile(mask, mult)\n",
        "\n",
        "\n",
        "def create_model():\n",
        "    encoder_inputs = keras.Input(shape=(None,), dtype=\"int32\", name=\"encoder_inputs\")\n",
        "    x = PositionalEmbedding(MAX_LENGTH, VOCAB_SIZE, EMBED_DIM)(encoder_inputs)\n",
        "    encoder_outputs = FNetEncoder(EMBED_DIM, LATENT_DIM)(x)\n",
        "    encoder = keras.Model(encoder_inputs, encoder_outputs)\n",
        "    decoder_inputs = keras.Input(shape=(None,), dtype=\"int32\", name=\"decoder_inputs\")\n",
        "    encoded_seq_inputs = keras.Input(\n",
        "        shape=(None, EMBED_DIM), name=\"decoder_state_inputs\"\n",
        "    )\n",
        "    x = PositionalEmbedding(MAX_LENGTH, VOCAB_SIZE, EMBED_DIM)(decoder_inputs)\n",
        "    x = FNetDecoder(EMBED_DIM, LATENT_DIM, NUM_HEADS)(x, encoded_seq_inputs)\n",
        "    x = layers.Dropout(0.5)(x)\n",
        "    decoder_outputs = layers.Dense(VOCAB_SIZE, activation=\"softmax\")(x)\n",
        "    decoder = keras.Model(\n",
        "        [decoder_inputs, encoded_seq_inputs], decoder_outputs, name=\"outputs\"\n",
        "    )\n",
        "    decoder_outputs = decoder([decoder_inputs, encoder_outputs])\n",
        "    fnet = keras.Model([encoder_inputs, decoder_inputs], decoder_outputs, name=\"fnet\")\n",
        "    return fnet"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mJ_B2fWtJMp-"
      },
      "source": [
        "fnet = create_model()\n",
        "fnet.compile(\"adam\", loss=\"sparse_categorical_crossentropy\", metrics=[\"accuracy\"])"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_X6P_Wk2JSBm",
        "outputId": "bb632755-9819-4c73-aa13-8e3178d517e1"
      },
      "source": [
        "fnet.summary()"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"fnet\"\n",
            "__________________________________________________________________________________________________\n",
            "Layer (type)                    Output Shape         Param #     Connected to                     \n",
            "==================================================================================================\n",
            "encoder_inputs (InputLayer)     [(None, None)]       0                                            \n",
            "__________________________________________________________________________________________________\n",
            "positional_embedding (Positiona (None, None, 256)    2107392     encoder_inputs[0][0]             \n",
            "__________________________________________________________________________________________________\n",
            "decoder_inputs (InputLayer)     [(None, None)]       0                                            \n",
            "__________________________________________________________________________________________________\n",
            "f_net_encoder (FNetEncoder)     (None, None, 256)    263936      positional_embedding[0][0]       \n",
            "__________________________________________________________________________________________________\n",
            "outputs (Functional)            (None, None, 8192)   8684288     decoder_inputs[0][0]             \n",
            "                                                                 f_net_encoder[0][0]              \n",
            "==================================================================================================\n",
            "Total params: 11,055,616\n",
            "Trainable params: 11,055,616\n",
            "Non-trainable params: 0\n",
            "__________________________________________________________________________________________________\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fiGAjRTtJW0d",
        "outputId": "ec7712b9-7173-4e77-8d30-33f50644e983"
      },
      "source": [
        "fnet.fit(train_dataset, epochs=1, validation_data=val_dataset)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "625/625 [==============================] - 94s 131ms/step - loss: 1.2642 - accuracy: 0.4545 - val_loss: 0.7286 - val_accuracy: 0.6731\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<keras.callbacks.History at 0x7fdbfb366710>"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "aSNSZGENJ4Kx",
        "outputId": "3f8d4c2b-0dad-4457-c7a8-47f1abffca35"
      },
      "source": [
        "VOCAB = vectorizer.get_vocabulary()\n",
        "\n",
        "\n",
        "def decode_sentence(input_sentence):\n",
        "    # Mapping the input sentence to tokens and adding start and end tokens\n",
        "    tokenized_input_sentence = vectorizer(\n",
        "        tf.constant(\"[start] \" + preprocess_text(input_sentence) + \" [end]\")\n",
        "    )\n",
        "    # Initializing the initial sentence consisting of only the start token.\n",
        "    tokenized_target_sentence = tf.expand_dims(VOCAB.index(\"[start]\"), 0)\n",
        "    decoded_sentence = \"\"\n",
        "\n",
        "    for i in range(MAX_LENGTH):\n",
        "        # Get the predictions\n",
        "        predictions = fnet.predict(\n",
        "            {\n",
        "                \"encoder_inputs\": tf.expand_dims(tokenized_input_sentence, 0),\n",
        "                \"decoder_inputs\": tf.expand_dims(\n",
        "                    tf.pad(\n",
        "                        tokenized_target_sentence,\n",
        "                        [[0, MAX_LENGTH - tf.shape(tokenized_target_sentence)[0]]],\n",
        "                    ),\n",
        "                    0,\n",
        "                ),\n",
        "            }\n",
        "        )\n",
        "        # Calculating the token with maximum probability and getting the corresponding word\n",
        "        sampled_token_index = tf.argmax(predictions[0, i, :])\n",
        "        sampled_token = VOCAB[sampled_token_index.numpy()]\n",
        "        # If sampled token is the end token then stop generating and return the sentence\n",
        "        if tf.equal(sampled_token_index, VOCAB.index(\"[end]\")):\n",
        "            break\n",
        "        decoded_sentence += sampled_token + \" \"\n",
        "        tokenized_target_sentence = tf.concat(\n",
        "            [tokenized_target_sentence, [sampled_token_index]], 0\n",
        "        )\n",
        "\n",
        "    return decoded_sentence\n",
        "\n",
        "\n",
        "decode_sentence(\"Where have you been all this time?\")"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'have you put some things have all one time ? mind '"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    }
  ]
}