{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "19ea4fe1",
   "metadata": {},
   "source": [
    "# Lesson 41: Encoder-decoder improvements activity\n",
    "\n",
    "In this activity, you will experiment with modifications to the encoder-decoder architecture from the lesson demo.\n",
    "\n",
    "1. **Embedding layer** - Add an embedding layer instead of one-hot encoding\n",
    "2. **Bidirectional encoder** - Make the encoder bidirectional\n",
    "3. **Hyperparameter tuning** - Experiment with latent dimension size\n",
    "\n",
    "## Notebook set-up\n",
    "\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e906ba3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.layers import Input, LSTM, Dense, Embedding, Bidirectional, Concatenate\n",
    "from tensorflow.keras.preprocessing.text import Tokenizer\n",
    "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
    "\n",
    "np.random.seed(315)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a09c40a0",
   "metadata": {},
   "source": [
    "## 1. Prepare data\n",
    "\n",
    "Using the same English-French translation pairs from the demo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46b5bdf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# English-French translation pairs\n",
    "pairs = [\n",
    "    ('hello', 'bonjour'),\n",
    "    ('goodbye', 'au revoir'),\n",
    "    ('thank you', 'merci'),\n",
    "    ('yes', 'oui'),\n",
    "    ('no', 'non'),\n",
    "    ('please', 's il vous plait'),\n",
    "    ('good morning', 'bonjour'),\n",
    "    ('good night', 'bonne nuit'),\n",
    "    ('how are you', 'comment allez vous'),\n",
    "    ('i am fine', 'je vais bien'),\n",
    "    ('what is your name', 'comment vous appelez vous'),\n",
    "    ('my name is', 'je m appelle'),\n",
    "    ('nice to meet you', 'enchanté'),\n",
    "    ('see you later', 'à bientôt'),\n",
    "    ('i love you', 'je t aime'),\n",
    "]\n",
    "\n",
    "# Prepare input/output texts with start/end tokens\n",
    "input_texts = [pair[0] for pair in pairs]\n",
    "target_texts = ['<start> ' + pair[1] + ' <end>' for pair in pairs]\n",
    "\n",
    "# Tokenize\n",
    "input_tokenizer = Tokenizer(filters='')\n",
    "input_tokenizer.fit_on_texts(input_texts)\n",
    "target_tokenizer = Tokenizer(filters='')\n",
    "target_tokenizer.fit_on_texts(target_texts)\n",
    "\n",
    "# Convert to sequences\n",
    "encoder_input_seq = input_tokenizer.texts_to_sequences(input_texts)\n",
    "decoder_input_seq = target_tokenizer.texts_to_sequences(target_texts)\n",
    "\n",
    "# Pad sequences\n",
    "max_encoder_len = max(len(seq) for seq in encoder_input_seq)\n",
    "max_decoder_len = max(len(seq) for seq in decoder_input_seq)\n",
    "\n",
    "encoder_input_data = pad_sequences(encoder_input_seq, maxlen=max_encoder_len, padding='post')\n",
    "decoder_input_data = pad_sequences(decoder_input_seq, maxlen=max_decoder_len, padding='post')\n",
    "\n",
    "# Decoder target (shifted by one position)\n",
    "decoder_target_data = np.zeros((len(pairs), max_decoder_len, len(target_tokenizer.word_index) + 1))\n",
    "for i, seq in enumerate(decoder_input_seq):\n",
    "    for t, word_idx in enumerate(seq[1:], start=0):  # Skip <start>\n",
    "        decoder_target_data[i, t, word_idx] = 1.0\n",
    "\n",
    "num_encoder_tokens = len(input_tokenizer.word_index) + 1\n",
    "num_decoder_tokens = len(target_tokenizer.word_index) + 1\n",
    "\n",
    "print(f'Encoder vocabulary size: {num_encoder_tokens}')\n",
    "print(f'Decoder vocabulary size: {num_decoder_tokens}')\n",
    "print(f'Max encoder sequence length: {max_encoder_len}')\n",
    "print(f'Max decoder sequence length: {max_decoder_len}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c50fa82",
   "metadata": {},
   "source": [
    "## 2. Baseline model (from demo)\n",
    "\n",
    "This is the basic encoder-decoder from the lesson."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9231489f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_baseline_model(latent_dim=64):\n",
    "    '''Build basic encoder-decoder with embedding layers.'''\n",
    "\n",
    "    # Encoder\n",
    "    encoder_inputs = Input(shape=(max_encoder_len,))\n",
    "    encoder_embedding = Embedding(num_encoder_tokens, latent_dim)(encoder_inputs)\n",
    "    encoder_lstm = LSTM(latent_dim, return_state=True)\n",
    "    _, state_h, state_c = encoder_lstm(encoder_embedding)\n",
    "    encoder_states = [state_h, state_c]\n",
    "    \n",
    "    # Decoder\n",
    "    decoder_inputs = Input(shape=(max_decoder_len,))\n",
    "    decoder_embedding = Embedding(num_decoder_tokens, latent_dim)(decoder_inputs)\n",
    "    decoder_lstm = LSTM(latent_dim, return_sequences=True, return_state=True)\n",
    "    decoder_outputs, _, _ = decoder_lstm(decoder_embedding, initial_state=encoder_states)\n",
    "    decoder_dense = Dense(num_decoder_tokens, activation='softmax')\n",
    "    decoder_outputs = decoder_dense(decoder_outputs)\n",
    "    \n",
    "    model = Model([encoder_inputs, decoder_inputs], decoder_outputs)\n",
    "    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])\n",
    "    \n",
    "    return model\n",
    "\n",
    "baseline_model = build_baseline_model()\n",
    "print(f'Baseline model parameters: {baseline_model.count_params():,}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f17f95e",
   "metadata": {},
   "source": [
    "## 3. Improved model with bidirectional encoder\n",
    "\n",
    "### Task 1: Add bidirectional encoder\n",
    "\n",
    "Modify the baseline model to use a bidirectional LSTM encoder. This allows the model to capture context from both directions.\n",
    "\n",
    "**Hints:**\n",
    "- Wrap the encoder LSTM with `Bidirectional`\n",
    "- Use `merge_mode='concat'` (default) to concatenate forward and backward states\n",
    "- The decoder latent dimension should be `2 * latent_dim` to match the concatenated states\n",
    "- Use `Concatenate()` to combine forward and backward states for `state_h` and `state_c`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3525ab0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_bidirectional_model(latent_dim=64):\n",
    "    '''Build encoder-decoder with bidirectional encoder.\n",
    "    \n",
    "    TODO: Implement the bidirectional encoder modifications.\n",
    "    '''\n",
    "\n",
    "    # Encoder\n",
    "    encoder_inputs = Input(shape=(max_encoder_len,))\n",
    "    encoder_embedding = Embedding(num_encoder_tokens, latent_dim)(encoder_inputs)\n",
    "    \n",
    "    # TODO: Replace LSTM with Bidirectional LSTM\n",
    "    # Hint: encoder_lstm = Bidirectional(LSTM(latent_dim, return_state=True))\n",
    "    encoder_lstm = None  # YOUR CODE HERE\n",
    "    \n",
    "    # TODO: Get outputs and states from bidirectional LSTM\n",
    "    # Hint: outputs, fwd_h, fwd_c, bwd_h, bwd_c = encoder_lstm(encoder_embedding)\n",
    "    # YOUR CODE HERE\n",
    "    \n",
    "    # TODO: Concatenate forward and backward states\n",
    "    # Hint: state_h = Concatenate()([fwd_h, bwd_h])\n",
    "    # Hint: state_c = Concatenate()([fwd_c, bwd_c])\n",
    "    encoder_states = None  # YOUR CODE HERE - should be [state_h, state_c]\n",
    "    \n",
    "    # Decoder (note: latent_dim * 2 because of bidirectional concatenation)\n",
    "    decoder_inputs = Input(shape=(max_decoder_len,))\n",
    "    decoder_embedding = Embedding(num_decoder_tokens, latent_dim)(decoder_inputs)\n",
    "    decoder_lstm = LSTM(latent_dim * 2, return_sequences=True, return_state=True)\n",
    "    decoder_outputs, _, _ = decoder_lstm(decoder_embedding, initial_state=encoder_states)\n",
    "    decoder_dense = Dense(num_decoder_tokens, activation='softmax')\n",
    "    decoder_outputs = decoder_dense(decoder_outputs)\n",
    "    \n",
    "    model = Model([encoder_inputs, decoder_inputs], decoder_outputs)\n",
    "    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "560d0e43",
   "metadata": {},
   "source": [
    "## 4. Train and compare models\n",
    "\n",
    "### Task 2: Compare model performance\n",
    "\n",
    "Train both models and compare their training accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "401d1dec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train baseline model\n",
    "print('Training baseline model...')\n",
    "baseline_model = build_baseline_model(latent_dim=64)\n",
    "baseline_history = baseline_model.fit(\n",
    "    [encoder_input_data, decoder_input_data],\n",
    "    decoder_target_data,\n",
    "    batch_size=4,\n",
    "    epochs=100,\n",
    "    verbose=0\n",
    ")\n",
    "print(f'Baseline final accuracy: {baseline_history.history[\"accuracy\"][-1]:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "426226a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Train bidirectional model (after completing Task 1)\n",
    "print('Training bidirectional model...')\n",
    "bidirectional_model = build_bidirectional_model(latent_dim=64)\n",
    "bidirectional_history = bidirectional_model.fit(\n",
    "    [encoder_input_data, decoder_input_data],\n",
    "    decoder_target_data,\n",
    "    batch_size=4,\n",
    "    epochs=100,\n",
    "    verbose=0\n",
    ")\n",
    "print(f'Bidirectional final accuracy: {bidirectional_history.history[\"accuracy\"][-1]:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9cd98f4",
   "metadata": {},
   "source": [
    "## 5. Hyperparameter experiment\n",
    "\n",
    "### Task 3: Experiment with latent dimension\n",
    "\n",
    "Try different values of `latent_dim` (e.g., 32, 64, 128, 256) and observe the effect on training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcb22b6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Experiment with different latent dimensions\n",
    "latent_dims = [32, 64, 128]  # Add more values to test\n",
    "results = {}\n",
    "\n",
    "for dim in latent_dims:\n",
    "    print(f'Training with latent_dim={dim}...')\n",
    "    model = build_baseline_model(latent_dim=dim)\n",
    "    history = model.fit(\n",
    "        [encoder_input_data, decoder_input_data],\n",
    "        decoder_target_data,\n",
    "        batch_size=4,\n",
    "        epochs=100,\n",
    "        verbose=0\n",
    "    )\n",
    "    results[dim] = {\n",
    "        'accuracy': history.history['accuracy'][-1],\n",
    "        'params': model.count_params()\n",
    "    }\n",
    "    print(f'  Accuracy: {results[dim][\"accuracy\"]:.4f}, Parameters: {results[dim][\"params\"]:,}')\n",
    "\n",
    "print('\\nSummary:')\n",
    "for dim, res in results.items():\n",
    "    print(f'latent_dim={dim:3d}: accuracy={res[\"accuracy\"]:.4f}, params={res[\"params\"]:,}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b77884f4",
   "metadata": {},
   "source": [
    "## 6. Analysis questions\n",
    "\n",
    "After completing the tasks above, answer these questions:\n",
    "\n",
    "1. How does the bidirectional encoder affect model performance? Why might this help for translation?\n",
    "2. What is the relationship between latent dimension and training accuracy?\n",
    "3. How does model size (number of parameters) scale with latent dimension?\n",
    "4. What other modifications could improve translation quality?"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}