transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer chatbot - implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensorflow version 2.10.0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "from __future__ import absolute_import, division, print_function, unicode_literals\n",
    "\n",
    "import sys\n",
    "\n",
    "import pickle\n",
    "import os\n",
    "import re\n",
    "import numpy as np\n",
    "from time import time\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "tf.keras.utils.set_random_seed(1234)\n",
    "\n",
    "print(f\"Tensorflow version {tf.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GPU Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]\n"
     ]
    }
   ],
   "source": [
    "print(tf.config.list_physical_devices())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensorflow uses GPU.\n"
     ]
    }
   ],
   "source": [
    "gpus = tf.config.list_physical_devices('GPU')\n",
    "if gpus:\n",
    "    print(\"Tensorflow uses GPU.\")\n",
    "else:\n",
    "    print(\"NOT GPU.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num GPUs Available:  1\n"
     ]
    }
   ],
   "source": [
    "print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)\n",
      "REPLICAS: 1\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    # TPU initialization (skipped if not available)\n",
    "    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n",
    "    print(\"Running on TPU {}\".format(tpu.cluster_spec().as_dict()[\"worker\"]))\n",
    "except ValueError:\n",
    "    tpu = None\n",
    "\n",
    "# Use GPU strategy if TPU is not available\n",
    "if tpu:\n",
    "    tf.config.experimental_connect_to_cluster(tpu)\n",
    "    tf.tpu.experimental.initialize_tpu_system(tpu)\n",
    "    strategy = tf.distribute.experimental.TPUStrategy(tpu)\n",
    "else:\n",
    "    # Use MirroredStrategy for GPU training\n",
    "    strategy = tf.distribute.MirroredStrategy()\n",
    "\n",
    "print(f\"REPLICAS: {strategy.num_replicas_in_sync}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Maximum sentence length\n",
    "MAX_LENGTH = 55\n",
    "\n",
    "# For tf.data.Dataset\n",
    "BATCH_SIZE = 64 * strategy.num_replicas_in_sync\n",
    "BUFFER_SIZE = 20000\n",
    "\n",
    "# For Transformer\n",
    "NUM_LAYERS = 3\n",
    "D_MODEL = 256\n",
    "NUM_HEADS = 8\n",
    "UNITS = 512\n",
    "DROPOUT = 0.1\n",
    "\n",
    "EPOCHS = 30"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load dataset\n",
    "The dataset consists of conversational pairs in Polish, previously processed to be ready for further action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "283389\n",
      "283389\n"
     ]
    }
   ],
   "source": [
    "def load_data(filename):\n",
    "    with open(filename, 'rb') as f:\n",
    "        return pickle.load(f)\n",
    "\n",
    "questions = load_data('pairs/questions-new.pkl')\n",
    "answers = load_data('pairs/answers-new.pkl')\n",
    "\n",
    "print(len(questions))\n",
    "print(len(answers))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pewnie zaraz będzie burza .\n",
      "ale gorąco , co ?\n"
     ]
    }
   ],
   "source": [
    "print(questions[2011])\n",
    "print(answers[2011])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tokenization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build tokenizer using tfds for both questions and answers\n",
    "tokenizer = tfds.deprecated.text.SubwordTextEncoder.build_from_corpus(\n",
    "    questions + answers, target_vocab_size=2**15\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.save_to_file('tokenizer_final_15_64')\n",
    "#tokenizer = tfds.deprecated.text.SubwordTextEncoder.load_from_file('tokenizer_new_15')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define start and end token to indicate the start and end of a sentence\n",
    "START_TOKEN, END_TOKEN = [tokenizer.vocab_size], [tokenizer.vocab_size + 1]\n",
    "\n",
    "# Vocabulary size plus start and end token\n",
    "VOCAB_SIZE = tokenizer.vocab_size + 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original: Cześć, jak się masz?\n",
      "Encoded: [32344, 510, 1073, 3790, 19, 7, 423, 32340]\n",
      "Decoded: Cześć, jak się masz?\n",
      "Original: Jestem chatbotem.\n",
      "Encoded: [32351, 20868, 47, 5775, 6790, 69, 32323]\n",
      "Decoded: Jestem chatbotem.\n"
     ]
    }
   ],
   "source": [
    "test_sentences = [\"Cześć, jak się masz?\", \"Jestem chatbotem.\"]\n",
    "for sentence in test_sentences:\n",
    "    encoded = tokenizer.encode(sentence)\n",
    "    decoded = tokenizer.decode(encoded)\n",
    "    print(f\"Original: {sentence}\")\n",
    "    print(f\"Encoded: {encoded}\")\n",
    "    print(f\"Decoded: {decoded}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tokenize, filter and pad sentences\n",
    "def tokenize_and_filter(inputs, outputs):\n",
    "    tokenized_inputs, tokenized_outputs = [], []\n",
    "\n",
    "    for (sentence1, sentence2) in zip(inputs, outputs):\n",
    "        # tokenize sentence\n",
    "        sentence1 = START_TOKEN + tokenizer.encode(sentence1) + END_TOKEN\n",
    "        sentence2 = START_TOKEN + tokenizer.encode(sentence2) + END_TOKEN\n",
    "        # check tokenized sentence max length\n",
    "        if len(sentence1) <= MAX_LENGTH and len(sentence2) <= MAX_LENGTH:\n",
    "            tokenized_inputs.append(sentence1)\n",
    "            tokenized_outputs.append(sentence2)\n",
    "\n",
    "    # pad tokenized sentences\n",
    "    tokenized_inputs = tf.keras.preprocessing.sequence.pad_sequences(\n",
    "        tokenized_inputs, maxlen=MAX_LENGTH, padding=\"post\"\n",
    "    )\n",
    "    tokenized_outputs = tf.keras.preprocessing.sequence.pad_sequences(\n",
    "        tokenized_outputs, maxlen=MAX_LENGTH, padding=\"post\"\n",
    "    )\n",
    "\n",
    "    return tokenized_inputs, tokenized_outputs\n",
    "\n",
    "\n",
    "questions, answers = tokenize_and_filter(questions, answers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vocab size: 32535\n",
      "Number of samples: 283356\n"
     ]
    }
   ],
   "source": [
    "print(f\"Vocab size: {VOCAB_SIZE}\")\n",
    "print(f\"Number of samples: {len(questions)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# decoder inputs use the previous target as input\n",
    "# remove START_TOKEN from targets\n",
    "dataset = tf.data.Dataset.from_tensor_slices(\n",
    "    (\n",
    "        {\"inputs\": questions, \"dec_inputs\": answers[:, :-1]},\n",
    "        {\"outputs\": answers[:, 1:]},\n",
    "    )\n",
    ")\n",
    "\n",
    "dataset = dataset.cache()\n",
    "dataset = dataset.shuffle(BUFFER_SIZE)\n",
    "dataset = dataset.batch(BATCH_SIZE)\n",
    "dataset = dataset.prefetch(tf.data.AUTOTUNE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset size: 4428\n"
     ]
    }
   ],
   "source": [
    "dataset_size = tf.data.experimental.cardinality(dataset).numpy()\n",
    "print(\"Dataset size:\", dataset_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Attention & MultiHead Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scaled_dot_product_attention(query, key, value, mask):\n",
    "    \"\"\"Calculate the attention weights.\"\"\"\n",
    "    matmul_qk = tf.matmul(query, key, transpose_b=True)\n",
    "\n",
    "    # scale matmul_qk\n",
    "    depth = tf.cast(tf.shape(key)[-1], tf.float32)\n",
    "    logits = matmul_qk / tf.math.sqrt(depth)\n",
    "\n",
    "    # add the mask to zero out padding tokens\n",
    "    if mask is not None:\n",
    "        logits += mask * -1e9\n",
    "\n",
    "    # softmax is normalized on the last axis (seq_len_k)\n",
    "    attention_weights = tf.nn.softmax(logits, axis=-1)\n",
    "\n",
    "    output = tf.matmul(attention_weights, value)\n",
    "\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttentionLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, d_model, num_heads, **kwargs):\n",
    "        assert d_model % num_heads == 0\n",
    "        super(MultiHeadAttentionLayer, self).__init__(**kwargs)\n",
    "        self.num_heads = num_heads\n",
    "        self.d_model = d_model\n",
    "\n",
    "        self.depth = d_model // self.num_heads\n",
    "\n",
    "        self.query_dense = tf.keras.layers.Dense(units=d_model)\n",
    "        self.key_dense = tf.keras.layers.Dense(units=d_model)\n",
    "        self.value_dense = tf.keras.layers.Dense(units=d_model)\n",
    "\n",
    "        self.dense = tf.keras.layers.Dense(units=d_model)\n",
    "\n",
    "    def get_config(self):\n",
    "        config = super(MultiHeadAttentionLayer, self).get_config()\n",
    "        config.update(\n",
    "            {\n",
    "                \"num_heads\": self.num_heads,\n",
    "                \"d_model\": self.d_model,\n",
    "            }\n",
    "        )\n",
    "        return config\n",
    "\n",
    "    def split_heads(self, inputs, batch_size):\n",
    "        inputs = tf.keras.layers.Lambda(\n",
    "            lambda inputs: tf.reshape(\n",
    "                inputs, shape=(batch_size, -1, self.num_heads, self.depth)\n",
    "            )\n",
    "        )(inputs)\n",
    "        return tf.keras.layers.Lambda(\n",
    "            lambda inputs: tf.transpose(inputs, perm=[0, 2, 1, 3])\n",
    "        )(inputs)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        query, key, value, mask = (\n",
    "            inputs[\"query\"],\n",
    "            inputs[\"key\"],\n",
    "            inputs[\"value\"],\n",
    "            inputs[\"mask\"],\n",
    "        )\n",
    "        batch_size = tf.shape(query)[0]\n",
    "\n",
    "        # linear layers\n",
    "        query = self.query_dense(query)\n",
    "        key = self.key_dense(key)\n",
    "        value = self.value_dense(value)\n",
    "\n",
    "        # split heads\n",
    "        query = self.split_heads(query, batch_size)\n",
    "        key = self.split_heads(key, batch_size)\n",
    "        value = self.split_heads(value, batch_size)\n",
    "\n",
    "        # scaled dot-product attention\n",
    "        scaled_attention = scaled_dot_product_attention(query, key, value, mask)\n",
    "        scaled_attention = tf.keras.layers.Lambda(\n",
    "            lambda scaled_attention: tf.transpose(scaled_attention, perm=[0, 2, 1, 3])\n",
    "        )(scaled_attention)\n",
    "\n",
    "        # concatenation of heads\n",
    "        concat_attention = tf.keras.layers.Lambda(\n",
    "            lambda scaled_attention: tf.reshape(\n",
    "                scaled_attention, (batch_size, -1, self.d_model)\n",
    "            )\n",
    "        )(scaled_attention)\n",
    "\n",
    "        # final linear layer\n",
    "        outputs = self.dense(concat_attention)\n",
    "\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_padding_mask(x):\n",
    "    mask = tf.cast(tf.math.equal(x, 0), tf.float32)\n",
    "    # (batch_size, 1, 1, sequence length)\n",
    "    return mask[:, tf.newaxis, tf.newaxis, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[0. 0. 1. 0. 1.]]]\n",
      "\n",
      "\n",
      " [[[1. 1. 1. 0. 0.]]]], shape=(2, 1, 1, 5), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "print(create_padding_mask(tf.constant([[1, 2, 0, 3, 0], [0, 0, 0, 4, 5]])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_look_ahead_mask(x):\n",
    "    seq_len = tf.shape(x)[1]\n",
    "    look_ahead_mask = 1 - tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)\n",
    "    padding_mask = create_padding_mask(x)\n",
    "    return tf.maximum(look_ahead_mask, padding_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[0. 1. 1. 1. 1.]\n",
      "   [0. 0. 1. 1. 1.]\n",
      "   [0. 0. 1. 1. 1.]\n",
      "   [0. 0. 1. 0. 1.]\n",
      "   [0. 0. 1. 0. 0.]]]], shape=(1, 1, 5, 5), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "print(create_look_ahead_mask(tf.constant([[1, 2, 0, 4, 5]])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Positional Encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(tf.keras.layers.Layer):\n",
    "    def __init__(self, position, d_model, **kwargs):\n",
    "        super(PositionalEncoding, self).__init__(**kwargs)\n",
    "        self.position = position\n",
    "        self.d_model = d_model\n",
    "        self.pos_encoding = self.positional_encoding(position, d_model)\n",
    "\n",
    "    def get_config(self):\n",
    "        config = super(PositionalEncoding, self).get_config()\n",
    "        config.update(\n",
    "            {\n",
    "                \"position\": self.position,\n",
    "                \"d_model\": self.d_model,\n",
    "            }\n",
    "        )\n",
    "        return config\n",
    "\n",
    "    def get_angles(self, position, i, d_model):\n",
    "        angles = 1 / tf.pow(10000, (2 * (i // 2)) / tf.cast(d_model, tf.float32))\n",
    "        return position * angles\n",
    "\n",
    "    def positional_encoding(self, position, d_model):\n",
    "        angle_rads = self.get_angles(\n",
    "            position=tf.range(position, dtype=tf.float32)[:, tf.newaxis],\n",
    "            i=tf.range(d_model, dtype=tf.float32)[tf.newaxis, :],\n",
    "            d_model=d_model,\n",
    "        )\n",
    "        # apply sin to even index in the array\n",
    "        sines = tf.math.sin(angle_rads[:, 0::2])\n",
    "        # apply cos to odd index in the array\n",
    "        cosines = tf.math.cos(angle_rads[:, 1::2])\n",
    "\n",
    "        pos_encoding = tf.concat([sines, cosines], axis=-1)\n",
    "        pos_encoding = pos_encoding[tf.newaxis, ...]\n",
    "        return tf.cast(pos_encoding, tf.float32)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        return inputs + self.pos_encoding[:, : tf.shape(inputs)[1], :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encoder_layer(units, d_model, num_heads, dropout, name=\"encoder_layer\"):\n",
    "    inputs = tf.keras.Input(shape=(None, d_model), name=\"inputs\")\n",
    "    padding_mask = tf.keras.Input(shape=(1, 1, None), name=\"padding_mask\")\n",
    "\n",
    "    attention = MultiHeadAttentionLayer(d_model, num_heads, name=\"attention\")(\n",
    "        {\"query\": inputs, \"key\": inputs, \"value\": inputs, \"mask\": padding_mask}\n",
    "    )\n",
    "    attention = tf.keras.layers.Dropout(rate=dropout)(attention)\n",
    "    add_attention = tf.keras.layers.add([inputs, attention])\n",
    "    attention = tf.keras.layers.LayerNormalization(epsilon=1e-6)(add_attention)\n",
    "\n",
    "    outputs = tf.keras.layers.Dense(units=units, activation=\"relu\")(attention)\n",
    "    outputs = tf.keras.layers.Dense(units=d_model)(outputs)\n",
    "    outputs = tf.keras.layers.Dropout(rate=dropout)(outputs)\n",
    "    add_attention = tf.keras.layers.add([attention, outputs])\n",
    "    outputs = tf.keras.layers.LayerNormalization(epsilon=1e-6)(add_attention)\n",
    "\n",
    "    return tf.keras.Model(inputs=[inputs, padding_mask], outputs=outputs, name=name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encoder(vocab_size, num_layers, units, d_model, num_heads, dropout, name=\"encoder\"):\n",
    "    inputs = tf.keras.Input(shape=(None,), name=\"inputs\")\n",
    "    padding_mask = tf.keras.Input(shape=(1, 1, None), name=\"padding_mask\")\n",
    "\n",
    "    embeddings = tf.keras.layers.Embedding(vocab_size, d_model)(inputs)\n",
    "    embeddings *= tf.keras.layers.Lambda(\n",
    "        lambda d_model: tf.math.sqrt(tf.cast(d_model, tf.float32))\n",
    "    )(d_model)\n",
    "    embeddings = PositionalEncoding(vocab_size, d_model)(embeddings)\n",
    "\n",
    "    outputs = tf.keras.layers.Dropout(rate=dropout)(embeddings)\n",
    "\n",
    "    for i in range(num_layers):\n",
    "        outputs = encoder_layer(\n",
    "            units=units,\n",
    "            d_model=d_model,\n",
    "            num_heads=num_heads,\n",
    "            dropout=dropout,\n",
    "            name=\"encoder_layer_{}\".format(i),\n",
    "        )([outputs, padding_mask])\n",
    "\n",
    "    return tf.keras.Model(inputs=[inputs, padding_mask], outputs=outputs, name=name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decoder_layer(units, d_model, num_heads, dropout, name=\"decoder_layer\"):\n",
    "    inputs = tf.keras.Input(shape=(None, d_model), name=\"inputs\")\n",
    "    enc_outputs = tf.keras.Input(shape=(None, d_model), name=\"encoder_outputs\")\n",
    "    look_ahead_mask = tf.keras.Input(shape=(1, None, None), name=\"look_ahead_mask\")\n",
    "    padding_mask = tf.keras.Input(shape=(1, 1, None), name=\"padding_mask\")\n",
    "\n",
    "    attention1 = MultiHeadAttentionLayer(d_model, num_heads, name=\"attention_1\")(\n",
    "        inputs={\n",
    "            \"query\": inputs,\n",
    "            \"key\": inputs,\n",
    "            \"value\": inputs,\n",
    "            \"mask\": look_ahead_mask,\n",
    "        }\n",
    "    )\n",
    "    add_attention = tf.keras.layers.add([attention1, inputs])\n",
    "    attention1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)(add_attention)\n",
    "\n",
    "    attention2 = MultiHeadAttentionLayer(d_model, num_heads, name=\"attention_2\")(\n",
    "        inputs={\n",
    "            \"query\": attention1,\n",
    "            \"key\": enc_outputs,\n",
    "            \"value\": enc_outputs,\n",
    "            \"mask\": padding_mask,\n",
    "        }\n",
    "    )\n",
    "    attention2 = tf.keras.layers.Dropout(rate=dropout)(attention2)\n",
    "    add_attention = tf.keras.layers.add([attention2, attention1])\n",
    "    attention2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)(add_attention)\n",
    "\n",
    "    outputs = tf.keras.layers.Dense(units=units, activation=\"relu\")(attention2)\n",
    "    outputs = tf.keras.layers.Dense(units=d_model)(outputs)\n",
    "    outputs = tf.keras.layers.Dropout(rate=dropout)(outputs)\n",
    "    add_attention = tf.keras.layers.add([outputs, attention2])\n",
    "    outputs = tf.keras.layers.LayerNormalization(epsilon=1e-6)(add_attention)\n",
    "\n",
    "    return tf.keras.Model(\n",
    "        inputs=[inputs, enc_outputs, look_ahead_mask, padding_mask],\n",
    "        outputs=outputs,\n",
    "        name=name,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decoder(vocab_size, num_layers, units, d_model, num_heads, dropout, name=\"decoder\"):\n",
    "    inputs = tf.keras.Input(shape=(None,), name=\"inputs\")\n",
    "    enc_outputs = tf.keras.Input(shape=(None, d_model), name=\"encoder_outputs\")\n",
    "    look_ahead_mask = tf.keras.Input(shape=(1, None, None), name=\"look_ahead_mask\")\n",
    "    padding_mask = tf.keras.Input(shape=(1, 1, None), name=\"padding_mask\")\n",
    "\n",
    "    embeddings = tf.keras.layers.Embedding(vocab_size, d_model)(inputs)\n",
    "    embeddings *= tf.keras.layers.Lambda(\n",
    "        lambda d_model: tf.math.sqrt(tf.cast(d_model, tf.float32))\n",
    "    )(d_model)\n",
    "    embeddings = PositionalEncoding(vocab_size, d_model)(embeddings)\n",
    "\n",
    "    outputs = tf.keras.layers.Dropout(rate=dropout)(embeddings)\n",
    "\n",
    "    for i in range(num_layers):\n",
    "        outputs = decoder_layer(\n",
    "            units=units,\n",
    "            d_model=d_model,\n",
    "            num_heads=num_heads,\n",
    "            dropout=dropout,\n",
    "            name=\"decoder_layer_{}\".format(i),\n",
    "        )(inputs=[outputs, enc_outputs, look_ahead_mask, padding_mask])\n",
    "\n",
    "    return tf.keras.Model(\n",
    "        inputs=[inputs, enc_outputs, look_ahead_mask, padding_mask],\n",
    "        outputs=outputs,\n",
    "        name=name,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transformer(\n",
    "    vocab_size, num_layers, units, d_model, num_heads, dropout, name=\"transformer\"\n",
    "):\n",
    "    inputs = tf.keras.Input(shape=(None,), name=\"inputs\")\n",
    "    dec_inputs = tf.keras.Input(shape=(None,), name=\"dec_inputs\")\n",
    "\n",
    "    enc_padding_mask = tf.keras.layers.Lambda(\n",
    "        create_padding_mask, output_shape=(1, 1, None), name=\"enc_padding_mask\"\n",
    "    )(inputs)\n",
    "    # mask the future tokens for decoder inputs at the 1st attention block\n",
    "    look_ahead_mask = tf.keras.layers.Lambda(\n",
    "        create_look_ahead_mask, output_shape=(1, None, None), name=\"look_ahead_mask\"\n",
    "    )(dec_inputs)\n",
    "    # mask the encoder outputs for the 2nd attention block\n",
    "    dec_padding_mask = tf.keras.layers.Lambda(\n",
    "        create_padding_mask, output_shape=(1, 1, None), name=\"dec_padding_mask\"\n",
    "    )(inputs)\n",
    "\n",
    "    enc_outputs = encoder(\n",
    "        vocab_size=vocab_size,\n",
    "        num_layers=num_layers,\n",
    "        units=units,\n",
    "        d_model=d_model,\n",
    "        num_heads=num_heads,\n",
    "        dropout=dropout,\n",
    "    )(inputs=[inputs, enc_padding_mask])\n",
    "\n",
    "    dec_outputs = decoder(\n",
    "        vocab_size=vocab_size,\n",
    "        num_layers=num_layers,\n",
    "        units=units,\n",
    "        d_model=d_model,\n",
    "        num_heads=num_heads,\n",
    "        dropout=dropout,\n",
    "    )(inputs=[dec_inputs, enc_outputs, look_ahead_mask, dec_padding_mask])\n",
    "\n",
    "    outputs = tf.keras.layers.Dense(units=vocab_size, name=\"outputs\")(dec_outputs)\n",
    "\n",
    "    return tf.keras.Model(inputs=[inputs, dec_inputs], outputs=outputs, name=name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function(y_true, y_pred):\n",
    "    y_true = tf.reshape(y_true, shape=(-1, MAX_LENGTH - 1))\n",
    "\n",
    "    loss = tf.keras.losses.SparseCategoricalCrossentropy(\n",
    "        from_logits=True, reduction=\"none\"\n",
    "    )(y_true, y_pred)\n",
    "\n",
    "    mask = tf.cast(tf.not_equal(y_true, 0), tf.float32)\n",
    "    loss = tf.multiply(loss, mask)\n",
    "\n",
    "    return tf.reduce_mean(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
    "    def __init__(self, d_model, warmup_steps=4000):\n",
    "        super(CustomSchedule, self).__init__()\n",
    "\n",
    "        self.d_model = tf.constant(d_model, dtype=tf.float32)\n",
    "        self.warmup_steps = warmup_steps\n",
    "\n",
    "    def get_config(self):\n",
    "        return {\"d_model\": self.d_model, \"warmup_steps\": self.warmup_steps}\n",
    "\n",
    "    def __call__(self, step):\n",
    "        arg1 = tf.math.rsqrt(step)\n",
    "        arg2 = step * (self.warmup_steps**-1.5)\n",
    "\n",
    "        return tf.math.multiply(\n",
    "            tf.math.rsqrt(self.d_model), tf.math.minimum(arg1, arg2)\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"transformer\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " inputs (InputLayer)            [(None, None)]       0           []                               \n",
      "                                                                                                  \n",
      " dec_inputs (InputLayer)        [(None, None)]       0           []                               \n",
      "                                                                                                  \n",
      " enc_padding_mask (Lambda)      (None, 1, 1, None)   0           ['inputs[0][0]']                 \n",
      "                                                                                                  \n",
      " encoder (Functional)           (None, None, 256)    9910272     ['inputs[0][0]',                 \n",
      "                                                                  'enc_padding_mask[0][0]']       \n",
      "                                                                                                  \n",
      " look_ahead_mask (Lambda)       (None, 1, None, Non  0           ['dec_inputs[0][0]']             \n",
      "                                e)                                                                \n",
      "                                                                                                  \n",
      " dec_padding_mask (Lambda)      (None, 1, 1, None)   0           ['inputs[0][0]']                 \n",
      "                                                                                                  \n",
      " decoder (Functional)           (None, None, 256)    10701312    ['dec_inputs[0][0]',             \n",
      "                                                                  'encoder[0][0]',                \n",
      "                                                                  'look_ahead_mask[0][0]',        \n",
      "                                                                  'dec_padding_mask[0][0]']       \n",
      "                                                                                                  \n",
      " outputs (Dense)                (None, None, 32535)  8361495     ['decoder[0][0]']                \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 28,973,079\n",
      "Trainable params: 28,973,079\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# clear backend\n",
    "tf.keras.backend.clear_session()\n",
    "\n",
    "learning_rate = CustomSchedule(D_MODEL)\n",
    "\n",
    "optimizer = tf.keras.optimizers.Adam(\n",
    "    learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9\n",
    ")\n",
    "\n",
    "\n",
    "def accuracy(y_true, y_pred):\n",
    "    # ensure labels have shape (batch_size, MAX_LENGTH - 1)\n",
    "    y_true = tf.reshape(y_true, shape=(-1, MAX_LENGTH - 1))\n",
    "    return tf.keras.metrics.sparse_categorical_accuracy(y_true, y_pred)\n",
    "\n",
    "\n",
    "# initialize and compile model within strategy scope\n",
    "model = transformer(\n",
    "    vocab_size=VOCAB_SIZE,\n",
    "    num_layers=NUM_LAYERS,\n",
    "    units=UNITS,\n",
    "    d_model=D_MODEL,\n",
    "    num_heads=NUM_HEADS,\n",
    "    dropout=DROPOUT,\n",
    ")\n",
    "\n",
    "model.compile(optimizer=optimizer, loss=loss_function, metrics=[accuracy])\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/30\n",
      "4428/4428 [==============================] - 1535s 347ms/step - loss: 0.6958 - accuracy: 0.0399\n",
      "Epoch 2/30\n",
      "4428/4428 [==============================] - 1238s 280ms/step - loss: 0.6653 - accuracy: 0.0424\n",
      "Epoch 3/30\n",
      "4428/4428 [==============================] - 1237s 279ms/step - loss: 0.6528 - accuracy: 0.0439\n",
      "Epoch 4/30\n",
      "4428/4428 [==============================] - 1241s 280ms/step - loss: 0.6484 - accuracy: 0.0449\n",
      "Epoch 5/30\n",
      "4428/4428 [==============================] - 1239s 280ms/step - loss: 0.6447 - accuracy: 0.0458\n",
      "Epoch 6/30\n",
      "4428/4428 [==============================] - 1239s 280ms/step - loss: 0.6414 - accuracy: 0.0465\n",
      "Epoch 7/30\n",
      "4428/4428 [==============================] - 1241s 280ms/step - loss: 0.6385 - accuracy: 0.0471\n",
      "Epoch 8/30\n",
      "4428/4428 [==============================] - 1240s 280ms/step - loss: 0.6362 - accuracy: 0.0477\n",
      "Epoch 9/30\n",
      "4428/4428 [==============================] - 1241s 280ms/step - loss: 0.6344 - accuracy: 0.0481\n",
      "Epoch 10/30\n",
      "4428/4428 [==============================] - 1243s 281ms/step - loss: 0.6333 - accuracy: 0.0486\n",
      "Epoch 11/30\n",
      "4428/4428 [==============================] - 1242s 281ms/step - loss: 0.6335 - accuracy: 0.0490\n",
      "Epoch 12/30\n",
      "4428/4428 [==============================] - 1241s 280ms/step - loss: 0.6351 - accuracy: 0.0493\n",
      "Epoch 13/30\n",
      "4428/4428 [==============================] - 1240s 280ms/step - loss: 0.6358 - accuracy: 0.0497\n",
      "Epoch 14/30\n",
      "4428/4428 [==============================] - 1241s 280ms/step - loss: 0.6348 - accuracy: 0.0500\n",
      "Epoch 15/30\n",
      "4428/4428 [==============================] - 1237s 279ms/step - loss: 0.6338 - accuracy: 0.0504\n",
      "Epoch 16/30\n",
      "4428/4428 [==============================] - 1237s 279ms/step - loss: 0.6327 - accuracy: 0.0507\n",
      "Epoch 17/30\n",
      "4428/4428 [==============================] - 1240s 280ms/step - loss: 0.6317 - accuracy: 0.0510\n",
      "Epoch 18/30\n",
      "4428/4428 [==============================] - 1240s 280ms/step - loss: 0.6306 - accuracy: 0.0513\n",
      "Epoch 19/30\n",
      "4428/4428 [==============================] - 1237s 279ms/step - loss: 0.6289 - accuracy: 0.0516\n",
      "Epoch 20/30\n",
      "4428/4428 [==============================] - 1238s 280ms/step - loss: 0.6272 - accuracy: 0.0519\n",
      "Epoch 21/30\n",
      "4428/4428 [==============================] - 1237s 279ms/step - loss: 0.6262 - accuracy: 0.0522\n",
      "Epoch 22/30\n",
      "4428/4428 [==============================] - 1236s 279ms/step - loss: 0.6249 - accuracy: 0.0525\n",
      "Epoch 23/30\n",
      "4428/4428 [==============================] - 1238s 280ms/step - loss: 0.6233 - accuracy: 0.0527\n",
      "Epoch 24/30\n",
      "4428/4428 [==============================] - 1237s 279ms/step - loss: 0.6222 - accuracy: 0.0529\n",
      "Epoch 25/30\n",
      "4428/4428 [==============================] - 1237s 279ms/step - loss: 0.6209 - accuracy: 0.0532\n",
      "Epoch 26/30\n",
      "4428/4428 [==============================] - 1238s 280ms/step - loss: 0.6198 - accuracy: 0.0534\n",
      "Epoch 27/30\n",
      "4428/4428 [==============================] - 1238s 280ms/step - loss: 0.6188 - accuracy: 0.0537\n",
      "Epoch 28/30\n",
      "4428/4428 [==============================] - 1238s 280ms/step - loss: 0.6172 - accuracy: 0.0539\n",
      "Epoch 29/30\n",
      "4428/4428 [==============================] - 1237s 279ms/step - loss: 0.6162 - accuracy: 0.0541\n",
      "Epoch 30/30\n",
      "4428/4428 [==============================] - 1239s 280ms/step - loss: 0.6149 - accuracy: 0.0544\n",
      "CPU times: total: 9h 15min 31s\n",
      "Wall time: 10h 24min 20s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x221547da820>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "model.fit(dataset, epochs=30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = \"./model/model-new-64.h5\"\n",
    "tf.keras.models.save_model(model, filepath=filename, include_optimizer=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = tf.keras.models.load_model(\n",
    "    \"./model/model-new-3.h5\",\n",
    "    custom_objects={\n",
    "        \"PositionalEncoding\": PositionalEncoding,\n",
    "        \"MultiHeadAttentionLayer\": MultiHeadAttentionLayer,\n",
    "    },\n",
    "    compile=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_sentence(sentence):\n",
    "    sentence = sentence.lower().strip()\n",
    "    # Creating a space between a word and the punctuation following it\n",
    "    sentence = re.sub(r\"([?.!,])\", r\" \\1 \", sentence)\n",
    "    sentence = re.sub(r'[\" \"]+', \" \", sentence)\n",
    "\n",
    "    # Replacing everything with space except (a-z, A-Z, \".\", \"?\", \"!\", \",\")\n",
    "    sentence = re.sub(r\"[^a-zA-ZąćęłńóśźżĄĆĘŁŃÓŚŹŻ0-9?.!,]+\", \" \", sentence)\n",
    "    sentence = sentence.strip()\n",
    "    return sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(sentence):\n",
    "    sentence = preprocess_sentence(sentence)\n",
    "\n",
    "    sentence = tf.expand_dims(\n",
    "        START_TOKEN + tokenizer.encode(sentence) + END_TOKEN, axis=0\n",
    "    )\n",
    "\n",
    "    output = tf.expand_dims(START_TOKEN, 0)\n",
    "\n",
    "    for i in range(MAX_LENGTH):\n",
    "        predictions = model(inputs=[sentence, output], training=False)\n",
    "\n",
    "        # select the last word from the seq_len dimension\n",
    "        predictions = predictions[:, -1:, :]\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 tf.equal(predicted_id, END_TOKEN[0]):\n",
    "            break\n",
    "\n",
    "        # concatenated 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)\n",
    "\n",
    "\n",
    "def predict(sentence):\n",
    "    prediction = evaluate(sentence)\n",
    "    predicted_sentence = tokenizer.decode(\n",
    "        [i for i in prediction if i < tokenizer.vocab_size]\n",
    "    )\n",
    "    return predicted_sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'nie wiem .'"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(\"wiesz co się stało?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input: Co u ciebie?\n",
      "Output: to jest to .\n",
      "\n",
      "Input: to jest to .\n",
      "Output: nie wiem , co to oznacza .\n",
      "\n",
      "Input: nie wiem , co to oznacza .\n",
      "Output: to nie jest śmieszne .\n",
      "\n",
      "Input: to nie jest śmieszne .\n",
      "Output: nie chcę się z tobą widywać .\n",
      "\n",
      "Input: nie chcę się z tobą widywać .\n",
      "Output: nie wiem , co to oznacza .\n",
      "\n"
     ]
    }
   ],
   "source": [
    "sentence = \"Co u ciebie?\"\n",
    "for _ in range(20):\n",
    "    print(f\"Input: {sentence}\")\n",
    "    sentence = predict(sentence)\n",
    "    print(f\"Output: {sentence}\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chatbot",
   "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}