{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1bd4b25-0369-48ac-b446-7a7acc0726a3",
   "metadata": {},
   "source": [
    "# Simple Transformer Model for Addition\n",
    "This notebook demonstrates how to build and train a simple transformer model to perform addition. We'll start by generating a dataset of simple addition equations and then train a transformer model on this dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "f4cadf99",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Requirement already satisfied: ipywidgets in /home/silentnova/.local/lib/python3.10/site-packages (8.1.1)\n",
      "Requirement already satisfied: comm>=0.1.3 in /home/silentnova/.local/lib/python3.10/site-packages (from ipywidgets) (0.2.1)\n",
      "Requirement already satisfied: ipython>=6.1.0 in /home/silentnova/.local/lib/python3.10/site-packages (from ipywidgets) (8.20.0)\n",
      "Requirement already satisfied: traitlets>=4.3.1 in /home/silentnova/.local/lib/python3.10/site-packages (from ipywidgets) (5.14.1)\n",
      "Requirement already satisfied: widgetsnbextension~=4.0.9 in /home/silentnova/.local/lib/python3.10/site-packages (from ipywidgets) (4.0.9)\n",
      "Requirement already satisfied: jupyterlab-widgets~=3.0.9 in /home/silentnova/.local/lib/python3.10/site-packages (from ipywidgets) (3.0.9)\n",
      "Requirement already satisfied: decorator in /home/silentnova/.local/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (5.1.1)\n",
      "Requirement already satisfied: jedi>=0.16 in /home/silentnova/.local/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (0.19.1)\n",
      "Requirement already satisfied: matplotlib-inline in /home/silentnova/.local/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (0.1.6)\n",
      "Requirement already satisfied: prompt-toolkit<3.1.0,>=3.0.41 in /home/silentnova/.local/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (3.0.43)\n",
      "Requirement already satisfied: pygments>=2.4.0 in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets) (2.17.2)\n",
      "Requirement already satisfied: stack-data in /home/silentnova/.local/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (0.6.3)\n",
      "Requirement already satisfied: exceptiongroup in /home/silentnova/.local/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (1.2.0)\n",
      "Requirement already satisfied: pexpect>4.3 in /usr/lib/python3/dist-packages (from ipython>=6.1.0->ipywidgets) (4.8.0)\n",
      "Requirement already satisfied: parso<0.9.0,>=0.8.3 in /home/silentnova/.local/lib/python3.10/site-packages (from jedi>=0.16->ipython>=6.1.0->ipywidgets) (0.8.3)\n",
      "Requirement already satisfied: wcwidth in /home/silentnova/.local/lib/python3.10/site-packages (from prompt-toolkit<3.1.0,>=3.0.41->ipython>=6.1.0->ipywidgets) (0.2.13)\n",
      "Requirement already satisfied: executing>=1.2.0 in /home/silentnova/.local/lib/python3.10/site-packages (from stack-data->ipython>=6.1.0->ipywidgets) (2.0.1)\n",
      "Requirement already satisfied: asttokens>=2.1.0 in /home/silentnova/.local/lib/python3.10/site-packages (from stack-data->ipython>=6.1.0->ipywidgets) (2.4.1)\n",
      "Requirement already satisfied: pure-eval in /home/silentnova/.local/lib/python3.10/site-packages (from stack-data->ipython>=6.1.0->ipywidgets) (0.2.2)\n",
      "Requirement already satisfied: six>=1.12.0 in /usr/lib/python3/dist-packages (from asttokens>=2.1.0->stack-data->ipython>=6.1.0->ipywidgets) (1.16.0)\n",
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Requirement already satisfied: tensorflow in /home/silentnova/.local/lib/python3.10/site-packages (2.15.0.post1)\n",
      "Requirement already satisfied: absl-py>=1.0.0 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (2.1.0)\n",
      "Requirement already satisfied: astunparse>=1.6.0 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (1.6.3)\n",
      "Requirement already satisfied: flatbuffers>=23.5.26 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (23.5.26)\n",
      "Requirement already satisfied: gast!=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (0.5.4)\n",
      "Requirement already satisfied: google-pasta>=0.1.1 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (0.2.0)\n",
      "Requirement already satisfied: h5py>=2.9.0 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (3.10.0)\n",
      "Requirement already satisfied: libclang>=13.0.0 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (16.0.6)\n",
      "Requirement already satisfied: ml-dtypes~=0.2.0 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (0.2.0)\n",
      "Requirement already satisfied: numpy<2.0.0,>=1.23.5 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (1.26.3)\n",
      "Requirement already satisfied: opt-einsum>=2.3.2 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (3.3.0)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from tensorflow) (23.2)\n",
      "Requirement already satisfied: protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.20.3 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (4.23.4)\n",
      "Requirement already satisfied: setuptools in /usr/lib/python3/dist-packages (from tensorflow) (59.6.0)\n",
      "Requirement already satisfied: six>=1.12.0 in /usr/lib/python3/dist-packages (from tensorflow) (1.16.0)\n",
      "Requirement already satisfied: termcolor>=1.1.0 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (2.4.0)\n",
      "Requirement already satisfied: typing-extensions>=3.6.6 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (4.9.0)\n",
      "Requirement already satisfied: wrapt<1.15,>=1.11.0 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (1.14.1)\n",
      "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (0.35.0)\n",
      "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (1.60.0)\n",
      "Requirement already satisfied: tensorboard<2.16,>=2.15 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (2.15.1)\n",
      "Requirement already satisfied: tensorflow-estimator<2.16,>=2.15.0 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (2.15.0)\n",
      "Requirement already satisfied: keras<2.16,>=2.15.0 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorflow) (2.15.0)\n",
      "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/lib/python3/dist-packages (from astunparse>=1.6.0->tensorflow) (0.37.1)\n",
      "Requirement already satisfied: google-auth<3,>=1.6.3 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorboard<2.16,>=2.15->tensorflow) (2.26.2)\n",
      "Requirement already satisfied: google-auth-oauthlib<2,>=0.5 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorboard<2.16,>=2.15->tensorflow) (1.2.0)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorboard<2.16,>=2.15->tensorflow) (3.5.2)\n",
      "Requirement already satisfied: requests<3,>=2.21.0 in /usr/lib/python3/dist-packages (from tensorboard<2.16,>=2.15->tensorflow) (2.25.1)\n",
      "Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorboard<2.16,>=2.15->tensorflow) (0.7.2)\n",
      "Requirement already satisfied: werkzeug>=1.0.1 in /home/silentnova/.local/lib/python3.10/site-packages (from tensorboard<2.16,>=2.15->tensorflow) (3.0.1)\n",
      "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /home/silentnova/.local/lib/python3.10/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow) (5.3.2)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /home/silentnova/.local/lib/python3.10/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow) (0.3.0)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /home/silentnova/.local/lib/python3.10/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow) (4.9)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /home/silentnova/.local/lib/python3.10/site-packages (from google-auth-oauthlib<2,>=0.5->tensorboard<2.16,>=2.15->tensorflow) (1.3.1)\n",
      "Requirement already satisfied: MarkupSafe>=2.1.1 in /home/silentnova/.local/lib/python3.10/site-packages (from werkzeug>=1.0.1->tensorboard<2.16,>=2.15->tensorflow) (2.1.4)\n",
      "Requirement already satisfied: pyasn1<0.6.0,>=0.4.6 in /home/silentnova/.local/lib/python3.10/site-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow) (0.5.1)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /usr/lib/python3/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<2,>=0.5->tensorboard<2.16,>=2.15->tensorflow) (3.2.0)\n",
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (1.26.3)\n"
     ]
    }
   ],
   "source": [
    "!pip install ipywidgets\n",
    "!pip install tensorflow\n",
    "!pip install numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e799a955-9d61-4809-9fd8-2aec69650d13",
   "metadata": {},
   "source": [
    "## Dataset Generation\n",
    "We will generate a dataset of simple addition equations. Each equation will be in the form of 'a + b = c' where 'a', 'b', and 'c' are integers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "9374dfb6-f072-44a2-a9d9-233efe87ac7f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input Equation: 12 + 14 =\n",
      "Answer: 26;\n",
      "\n",
      "Input Equation: 28 + 73 =\n",
      "Answer: 101;\n",
      "\n",
      "Input Equation: 51 + 31 =\n",
      "Answer: 82;\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "from tensorflow.keras.preprocessing.text import Tokenizer\n",
    "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
    "\n",
    "def generate_addition_data(num_samples, stop_token=';'):\n",
    "    input_equations = []\n",
    "    answers = []\n",
    "    for _ in range(num_samples):\n",
    "        a = random.randint(0, 99)\n",
    "        b = random.randint(0, 99)\n",
    "        input_eq = f\"{a} + {b} =\"\n",
    "        answer = str(a + b) + stop_token  # Append the stop token to each answer\n",
    "        input_equations.append(input_eq)\n",
    "        answers.append(answer)\n",
    "    return input_equations, answers\n",
    "\n",
    "num_samples = 500000\n",
    "input_equations, answers = generate_addition_data(num_samples)\n",
    "\n",
    "\n",
    "num_samples_to_print = 3\n",
    "for i in range(num_samples_to_print):\n",
    "    print(f\"Input Equation: {input_equations[i]}\")\n",
    "    print(f\"Answer: {answers[i]}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95d98ddf-ca76-4139-ab88-af21659f8f7a",
   "metadata": {},
   "source": [
    "## Data Preprocessing\n",
    "We'll convert the equations into a format suitable for training a transformer model. This includes tokenization and converting tokens to numerical format.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "477da248-a90c-4a1d-b307-8a786aab9750",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input Equation: 12 + 14 =\n",
      "Tokenized Input Sequence: [2, 9, 1, 3, 1, 2, 12, 1, 4]\n",
      "Padded Input Sequence: [ 2  9  1  3  1  2 12  1  4]\n",
      "Answer: 26;\n",
      "Tokenized Answer Sequence: [9, 10, 5]\n",
      "Padded Answer Sequence: [ 9 10  5  0  0  0  0  0  0]\n",
      "\n",
      "Input Equation: 28 + 73 =\n",
      "Tokenized Input Sequence: [9, 8, 1, 3, 1, 7, 6, 1, 4]\n",
      "Padded Input Sequence: [9 8 1 3 1 7 6 1 4]\n",
      "Answer: 101;\n",
      "Tokenized Answer Sequence: [2, 14, 2, 5]\n",
      "Padded Answer Sequence: [ 2 14  2  5  0  0  0  0  0]\n",
      "\n",
      "Input Equation: 51 + 31 =\n",
      "Tokenized Input Sequence: [11, 2, 1, 3, 1, 6, 2, 1, 4]\n",
      "Padded Input Sequence: [11  2  1  3  1  6  2  1  4]\n",
      "Answer: 82;\n",
      "Tokenized Answer Sequence: [8, 9, 5]\n",
      "Padded Answer Sequence: [8 9 5 0 0 0 0 0 0]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Tokenization\n",
    "tokenizer = Tokenizer(char_level=True)  # Adjust tokenizer settings as needed\n",
    "tokenizer.fit_on_texts(input_equations + answers)\n",
    "\n",
    "# Convert to sequences\n",
    "input_sequences = tokenizer.texts_to_sequences(input_equations)\n",
    "answer_sequences = tokenizer.texts_to_sequences(answers)\n",
    "\n",
    "# Padding sequences\n",
    "max_len_input = max([len(seq) for seq in input_sequences])\n",
    "max_len_answer = max([len(seq) for seq in answer_sequences])\n",
    "max_len = max(max_len_input, max_len_answer)\n",
    "\n",
    "input_sequences_padded = pad_sequences(input_sequences, maxlen=max_len, padding='post')\n",
    "answer_sequences_padded = pad_sequences(answer_sequences, maxlen=max_len, padding='post')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "for i in range(num_samples_to_print):\n",
    "    print(f\"Input Equation: {input_equations[i]}\")\n",
    "    print(f\"Tokenized Input Sequence: {input_sequences[i]}\")\n",
    "    print(f\"Padded Input Sequence: {input_sequences_padded[i]}\")\n",
    "    \n",
    "    print(f\"Answer: {answers[i]}\")\n",
    "    print(f\"Tokenized Answer Sequence: {answer_sequences[i]}\")\n",
    "    print(f\"Padded Answer Sequence: {answer_sequences_padded[i]}\")\n",
    "    \n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa090f72",
   "metadata": {},
   "source": [
    "The tokenized input sequences you see are token IDs generated by Keras' Tokenizer. Let me explain how this works:\n",
    "\n",
    " - Tokenizer Creation: When we create a Tokenizer instance with char_level=True, the tokenizer treats each unique character as a distinct token.\n",
    "\n",
    " - Fitting the Tokenizer: By calling tokenizer.fit_on_texts(input_equations + answers), we are essentially instructing the tokenizer to go through all the characters in our dataset (both the input equations and the answers) and assign a unique integer ID to each different character.\n",
    "\n",
    " - Token IDs: The tokenizer then creates a mapping from characters to these integer IDs. For example, it might assign '1' to '+', '2' to '=', '3' to '0', '4' to ';', and so on for all unique characters (including all digits from '0' to '9'). The exact mapping depends on the order in which the characters are encountered and their frequency.\n",
    "\n",
    " - Tokenization Process: When we convert the text data into sequences using tokenizer.texts_to_sequences(...), each character in the input is replaced by its corresponding integer ID based on the mapping created by the tokenizer.\n",
    "\n",
    "For instance, if the input equation is \"49 + 51 =\", and the tokenizer has assigned '6' to '4', '13' to '9', '1' to ' ', '3' to '+', '12' to '5', and '2' to '=', then the tokenized input sequence for this equation would be [6, 13, 1, 3, 1, 12, 2, 1, 4].\n",
    "\n",
    "These token IDs are used throughout the model for processing, and they are crucial for both understanding the input data and generating predictions. The model learns to associate these tokens with their meaning in the context of addition operations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d89bad6-39a0-49ec-8754-931269ea4701",
   "metadata": {},
   "source": [
    "## Building the Transformer Model\n",
    "We will define a simple transformer model suitable for our task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "00811e2d-2d60-49be-8c38-e88bf6ecbdfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.layers import Input, Embedding, MultiHeadAttention, LayerNormalization, Dropout, Dense\n",
    "from tensorflow.keras.layers import GlobalAveragePooling1D\n",
    "from tensorflow.keras.layers import Masking\n",
    "import tensorflow as tf\n",
    "\n",
    "# Transformer block as a custom layer\n",
    "class TransformerBlock(tf.keras.layers.Layer):\n",
    "    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):\n",
    "        super(TransformerBlock, self).__init__()\n",
    "        self.att = MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)\n",
    "        self.ffn = tf.keras.Sequential(\n",
    "            [Dense(ff_dim, activation=\"relu\"), Dense(embed_dim),]\n",
    "        )\n",
    "        self.layernorm1 = LayerNormalization(epsilon=1e-6)\n",
    "        self.layernorm2 = LayerNormalization(epsilon=1e-6)\n",
    "        self.dropout1 = Dropout(rate)\n",
    "        self.dropout2 = Dropout(rate)\n",
    "\n",
    "    def call(self, inputs, training):\n",
    "        attn_output = self.att(inputs, inputs)\n",
    "        attn_output = self.dropout1(attn_output, training=training)\n",
    "        out1 = self.layernorm1(inputs + attn_output)\n",
    "        ffn_output = self.ffn(out1)\n",
    "        ffn_output = self.dropout2(ffn_output, training=training)\n",
    "        return self.layernorm2(out1 + ffn_output)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "4c69b1b4-070c-4cdc-8f33-ceb7f8fcaa63",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_3\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " input_4 (InputLayer)        [(None, 9)]               0         \n",
      "                                                                 \n",
      " embedding_3 (Embedding)     (None, 9, 64)             960       \n",
      "                                                                 \n",
      " masking (Masking)           (None, 9, 64)             0         \n",
      "                                                                 \n",
      " transformer_block_3 (Trans  (None, 9, 64)             446336    \n",
      " formerBlock)                                                    \n",
      "                                                                 \n",
      " global_average_pooling1d_3  (None, 64)                0         \n",
      "  (GlobalAveragePooling1D)                                       \n",
      "                                                                 \n",
      " dense_11 (Dense)            (None, 15)                975       \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 448271 (1.71 MB)\n",
      "Trainable params: 448271 (1.71 MB)\n",
      "Non-trainable params: 0 (0.00 Byte)\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "vocab_size = len(tokenizer.word_index) + 1\n",
    "embedding_dim = 64\n",
    "num_heads = 2\n",
    "ff_dim = 3200\n",
    "\n",
    "\n",
    "inputs = Input(shape=(max_len,))\n",
    "x = Embedding(vocab_size, embedding_dim)(inputs)\n",
    "x = Masking(mask_value=0)(x)  # Add Masking layer after Embedding\n",
    "x = TransformerBlock(embedding_dim, num_heads, ff_dim)(x)\n",
    "x = GlobalAveragePooling1D()(x)\n",
    "outputs = Dense(vocab_size, activation='softmax')(x)\n",
    "\n",
    "model = Model(inputs=inputs, outputs=outputs)\n",
    "model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f08462c-97c9-4588-80d1-80f818fe0269",
   "metadata": {},
   "source": [
    "## Model Training\n",
    "We will now train the transformer model on our generated dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "913ac294-5117-4b7f-842d-7af2c1a437f7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_input shape: (400000, 9)\n",
      "val_input shape: (100000, 9)\n",
      "train_answers_one_hot shape: (400000, 15)\n",
      "val_answers_one_hot shape: (100000, 15)\n",
      "Epoch 1/10\n",
      "12500/12500 [==============================] - 102s 8ms/step - loss: 0.9361 - accuracy: 0.6122 - val_loss: 0.8255 - val_accuracy: 0.6275\n",
      "Epoch 2/10\n",
      "12500/12500 [==============================] - 104s 8ms/step - loss: 0.8035 - accuracy: 0.6288 - val_loss: 0.7792 - val_accuracy: 0.6329\n",
      "Epoch 3/10\n",
      "12500/12500 [==============================] - 104s 8ms/step - loss: 0.7768 - accuracy: 0.6330 - val_loss: 0.7589 - val_accuracy: 0.6299\n",
      "Epoch 4/10\n",
      "12500/12500 [==============================] - 104s 8ms/step - loss: 0.7650 - accuracy: 0.6339 - val_loss: 0.7563 - val_accuracy: 0.6328\n",
      "Epoch 5/10\n",
      "12500/12500 [==============================] - 96s 8ms/step - loss: 0.7587 - accuracy: 0.6343 - val_loss: 0.7631 - val_accuracy: 0.6364\n",
      "Epoch 6/10\n",
      "12500/12500 [==============================] - 15804s 1s/step - loss: 0.7546 - accuracy: 0.6360 - val_loss: 0.7520 - val_accuracy: 0.6338\n",
      "Epoch 7/10\n",
      "12500/12500 [==============================] - 117s 9ms/step - loss: 0.7523 - accuracy: 0.6359 - val_loss: 0.7477 - val_accuracy: 0.6353\n",
      "Epoch 8/10\n",
      "12500/12500 [==============================] - 118s 9ms/step - loss: 0.7500 - accuracy: 0.6359 - val_loss: 0.7687 - val_accuracy: 0.6251\n",
      "Epoch 9/10\n",
      "12500/12500 [==============================] - 109s 9ms/step - loss: 0.7487 - accuracy: 0.6364 - val_loss: 0.7486 - val_accuracy: 0.6348\n",
      "Epoch 10/10\n",
      "12500/12500 [==============================] - 89s 7ms/step - loss: 0.7477 - accuracy: 0.6360 - val_loss: 0.7479 - val_accuracy: 0.6314\n"
     ]
    }
   ],
   "source": [
    "# Splitting dataset into training and validation\n",
    "train_size = int(0.8 * len(input_sequences_padded))\n",
    "train_input = input_sequences_padded[:train_size]\n",
    "train_answers = answer_sequences_padded[:train_size]\n",
    "\n",
    "\n",
    "val_input = input_sequences_padded[train_size:]\n",
    "val_answers = answer_sequences_padded[train_size:]\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "\n",
    "# Assuming each entry in your target data is an integer class label\n",
    "num_classes = 15  # as per your model's output\n",
    "\n",
    "\n",
    "\n",
    "# Flatten the target data\n",
    "train_answers_flattened = train_answers[:, 0]  # Assuming the class label is in the first column\n",
    "val_answers_flattened = val_answers[:, 0]\n",
    "\n",
    "# Apply one-hot encoding\n",
    "train_answers_one_hot = to_categorical(train_answers_flattened, num_classes=num_classes)\n",
    "val_answers_one_hot = to_categorical(val_answers_flattened, num_classes=num_classes)\n",
    "\n",
    "\n",
    "\n",
    "print(\"train_input shape:\", train_input.shape)\n",
    "print(\"val_input shape:\", val_input.shape)\n",
    "print(\"train_answers_one_hot shape:\", train_answers_one_hot.shape)\n",
    "print(\"val_answers_one_hot shape:\", val_answers_one_hot.shape)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# Training\n",
    "epochs = 10  # Adjust as needed\n",
    "model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])\n",
    "\n",
    "# Then continue with training\n",
    "history = model.fit(train_input, train_answers_one_hot, validation_data=(val_input, val_answers_one_hot), epochs=epochs, batch_size=32)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc146793-29e8-4bda-9423-197d05cc12f7",
   "metadata": {},
   "source": [
    "## Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "353ef3f0-3e85-4934-99ab-4a108f3e3cc5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - 0s 87ms/step\n",
      "1/1 [==============================] - 0s 12ms/step\n",
      "1/1 [==============================] - 0s 12ms/step\n",
      "1/1 [==============================] - 0s 11ms/step\n",
      "1/1 [==============================] - 0s 12ms/step\n",
      "1/1 [==============================] - 0s 12ms/step\n",
      "1/1 [==============================] - 0s 11ms/step\n",
      "1/1 [==============================] - 0s 11ms/step\n",
      "1/1 [==============================] - 0s 11ms/step\n",
      "1/1 [==============================] - 0s 11ms/step\n",
      "Predicted result: 9114461135\n"
     ]
    }
   ],
   "source": [
    "def predict_until_stop(model, tokenizer, input_text, max_length=10, stop_token=';'):\n",
    "    # Tokenizing the input\n",
    "    input_seq = tokenizer.texts_to_sequences([input_text])\n",
    "    # Padding the sequence\n",
    "    input_padded = pad_sequences(input_seq, maxlen=max_len, padding='post')\n",
    "\n",
    "    predicted_sequence = []\n",
    "    for _ in range(max_length):\n",
    "        # Make a prediction\n",
    "        prediction = model.predict(input_padded)\n",
    "        predicted_token_index = np.argmax(prediction, axis=1)[0]\n",
    "        predicted_token = tokenizer.index_word[predicted_token_index]\n",
    "\n",
    "        # Append to the sequence and break if stop token is predicted\n",
    "        if predicted_token == stop_token:\n",
    "            break\n",
    "        predicted_sequence.append(predicted_token)\n",
    "\n",
    "        # Update the input by shifting left and adding the new token at the end\n",
    "        input_padded = np.roll(input_padded, -1, axis=1)\n",
    "        input_padded[0, -1] = predicted_token_index\n",
    "\n",
    "    return ''.join(predicted_sequence)\n",
    "\n",
    "# Example usage\n",
    "new_input = \"73 + 22 =\"\n",
    "predicted_result = predict_until_stop(model, tokenizer, new_input, max_length=10)\n",
    "print(\"Predicted result:\", predicted_result)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}