visualize.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f30d3ba5-5d15-4a37-8483-1592fcd6db89",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from bertviz import head_view, model_view\n",
    "\n",
    "from transformer_implementation import Transformer, Tokenizer, TransformerConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe99d1ac-451d-45c8-9a5d-28799680a1b9",
   "metadata": {},
   "source": [
    "## Init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9a1014fd-c546-4959-b9b6-996af4f2dd88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# init tokenizer\n",
    "tokenizer = Tokenizer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4c4a4183-8126-4b94-8130-7541b0ca4d50",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TransformerConfig:\n",
      "\n",
      "Tokenizer:\n",
      "+---------------------+\n",
      "| vocab_size : 100277 |\n",
      "| BOS_IDX    : 100264 |\n",
      "| EOS_IDX    : 100265 |\n",
      "| PAD_IDX    : 100266 |\n",
      "+---------------------+\n",
      "\n",
      "Data:\n",
      "+---------------------------------+\n",
      "| block_size              : 64    |\n",
      "| batch_size              : 12    |\n",
      "| train_data_size         : 30000 |\n",
      "| grad_accumulation_steps : 40    |\n",
      "+---------------------------------+\n",
      "\n",
      "Model:\n",
      "+-----------------+\n",
      "| n_layer : 6     |\n",
      "| n_head  : 8     |\n",
      "| n_embd  : 512   |\n",
      "| dropout : 0.1   |\n",
      "| bias    : 0     |\n",
      "+-----------------+\n",
      "\n",
      "Training Loop:\n",
      "+-------------------+\n",
      "| max_epochs : 150  |\n",
      "| max_iters  : 2499 |\n",
      "| eval_iters : 249  |\n",
      "+-------------------+\n",
      "\n",
      "AdamW Optimizer:\n",
      "+------------------------+\n",
      "| learning_rate : 0.0006 |\n",
      "| beta1         : 0.9    |\n",
      "| beta2         : 0.95   |\n",
      "| weight_decay  : 0.1    |\n",
      "| eps           : 1e-09  |\n",
      "+------------------------+\n",
      "\n",
      "System:\n",
      "+-----------------------+\n",
      "| device_type : cuda    |\n",
      "| device      : cuda    |\n",
      "| dtype       : float16 |\n",
      "| compile     : 1       |\n",
      "+-----------------------+\n",
      "\n",
      "DDP Settings:\n",
      "+-----------------+\n",
      "| backend : nccl  |\n",
      "| ddp     : 0     |\n",
      "+-----------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# init config\n",
    "config = TransformerConfig(\n",
    "    vocab_size = tokenizer.vocab_size(),\n",
    "    block_size = 64,\n",
    "    n_layer = 6, # 6,\n",
    "    n_head = 8,\n",
    "    n_embd = 512,\n",
    "    max_epochs=150,\n",
    "    train_data_size = 30000, # batch * 500 iters\n",
    "    max_iters = int(30000/12)-1, # dataset_size / batch_size\n",
    "    eval_iters = int(3000/12)-1, # val_dataset_size / batch_size\n",
    "    BOS_IDX = tokenizer.BOS_IDX,\n",
    "    EOS_IDX = tokenizer.EOS_IDX,\n",
    "    PAD_IDX = tokenizer.PAD_IDX,\n",
    ")\n",
    "print(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0766891a-8164-45b0-8b7e-ce76eebaa75c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of Encoder parameters: 70.22M\n",
      "number of Decoder parameters: 76.52M\n",
      "Total number of parameters: 146.74M\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create model\n",
    "model = Transformer(config)\n",
    "# model.load_model(\"./out/transformer-train2.pth\")\n",
    "model.load_model(\"./out/train_model.pth\")\n",
    "model = model.to(config.device)\n",
    "model.eval()\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1e28e8b8-a6b7-4591-9ac3-ffc8770ed362",
   "metadata": {},
   "outputs": [],
   "source": [
    "def translate(sentences, tokenizer, model, config):\n",
    "    \"\"\"\n",
    "    This function tokenizes input sentences, translates them using the provided model,\n",
    "    and decodes the output into human-readable text. It also returns the attention dictionary from the model.\n",
    "\n",
    "    Args:\n",
    "        - sentences (list[str]): List of sentences to be translated.\n",
    "        - tokenizer (Tokenizer): Tokenizer used for encoding and decoding sequences.\n",
    "        - model (Transformer): The model used for translation.\n",
    "        - config (Config): The configuration object that defines parameters like block_size.\n",
    "\n",
    "    Returns:\n",
    "        - decode_output (list[str]): List of translated sentences.\n",
    "        - attn (dict): Dictionary containing attention information from the last layer of the model.\n",
    "    \"\"\"\n",
    "    # Tokenize sentences\n",
    "    tknzr = tokenizer.encoder\n",
    "    sequences = []\n",
    "    masks =  []\n",
    "\n",
    "    # Encode each sentence and add it to the list of sequences\n",
    "    for sentence in sentences:\n",
    "        sequence = tokenizer.sequence_padding(tknzr.encode(sentence), config.block_size).unsqueeze(dim=0)\n",
    "        mask = tokenizer.generate_padding_mask(sequence, device = config.device)\n",
    "        sequences.append(sequence)\n",
    "        masks.append(mask)\n",
    "\n",
    "    # Concatenate the sequences into a tensor\n",
    "    sequences = torch.cat(sequences, dim=0)\n",
    "    masks = torch.cat(masks, dim=0).to(config.device)\n",
    "\n",
    "    # Set the model to evaluation mode and translate sentences\n",
    "    model.eval()\n",
    "    outputs, attn = model.generate(\n",
    "        src=sequences.to(config.device),\n",
    "        idx=torch.full((1, 1), tokenizer.BOS_IDX).long().to(config.device),\n",
    "        src_mask=masks,\n",
    "        temperature = 0.5,\n",
    "        max_new_tokens=config.block_size,\n",
    "        top_k= 200,\n",
    "    )\n",
    "    # Initialize a list to store the decoded sentences\n",
    "    decode_output = []\n",
    "    # Decode each output sequence and add it to the list of decoded outputs\n",
    "    for output in outputs:\n",
    "        output = tokenizer.sequence_cleaner(output)\n",
    "        decode_output += [tknzr.decode(output)]\n",
    "\n",
    "    # Return the decoded sentences and the attention dictionary\n",
    "    return decode_output, attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a82890e5-6545-4d99-aac5-c687324c8dd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# input = [\"\"\"This is an issue which, as you will understand, will have to be addressed at IGC level and is an issue that only a Special European Council can authorise the Portuguese Presidency to pursue, because this kind of issue falls outside the institutional framework in which the forthcoming Intergovernmental Conference will take place.\"\"\"]\n",
    "# input = [\"\"\"This is a question that interests the government\"\"\"]\n",
    "input = [\"\"\"I am the French president\"\"\"]\n",
    "# expected_output = ['Je suis un professeur.']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5d52ba69-9237-452e-923e-2749af2125ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs, attentions = translate(input, tokenizer, model, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c704cd98-63d0-48cd-809c-18dd179be8ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Je suis le président français.']"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3a3d80d0-2107-4e6c-b8b6-c49012308c60",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_attn(input, output, attentions, batch: int = 0):\n",
    "    \"\"\"\n",
    "    This function formats the attention outputs and tokenized inputs and outputs for easier interpretation and visualization.\n",
    "\n",
    "    Args:\n",
    "        - input (str): The original input sentence.\n",
    "        - output (str): The translated output sentence.\n",
    "        - attentions (dict): A dictionary containing the attention information from the model.\n",
    "        - batch (int, optional): The batch index to format. Defaults to 0.\n",
    "\n",
    "    Returns:\n",
    "        - tokens_input (list[str]): The tokenized input sentence, padded to max_len.\n",
    "        - tokens_output (list[str]): The tokenized output sentence, padded to max_len.\n",
    "        - tensor_encoder_attn (torch.Tensor): The attention tensor for the encoder, trimmed and reshaped.\n",
    "        - tensor_cross_attn (torch.Tensor): The cross-attention tensor, trimmed and reshaped.\n",
    "        - tensor_decoder_attn (torch.Tensor): The attention tensor for the decoder, trimmed and reshaped.\n",
    "    \"\"\"\n",
    "\n",
    "    # Stack the attention tensors along a new dimension\n",
    "    tensor_encoder_attn = torch.stack(attentions['encoder_attn'], dim=0)\n",
    "    tensor_cross_attn = torch.stack(attentions['cross_attn'], dim=0)\n",
    "    tensor_decoder_attn = torch.stack(attentions['decoder_attn'], dim=0)\n",
    "\n",
    "    # Tokenize the input and output sentences\n",
    "    tokens_input = tokenizer.tokenize_from_str(input[batch])\n",
    "    tokens_output = tokenizer.tokenize_from_str(output[batch])\n",
    "\n",
    "    # Find the maximum length of the input and output tokens\n",
    "    max_len = max(len(tokens_input), len(tokens_output))\n",
    "    # If the input tokens are shorter than the max length, pad with empty strings\n",
    "    if len(tokens_input) < max_len:\n",
    "        tokens_input = tokens_input + [''] * (max_len - len(tokens_input))\n",
    "    # Otherwise, pad the output tokens with empty strings\n",
    "    else:\n",
    "        tokens_output = tokens_output + [''] * (max_len - len(tokens_output))\n",
    "\n",
    "    # Trim and reshape the attention tensors\n",
    "    tensor_encoder_attn = tensor_encoder_attn[:, batch:batch+1, :, 0:max_len, 0:max_len] # layers, batch, heads, seq_len, seq_len\n",
    "    tensor_cross_attn = tensor_cross_attn[:, batch:batch+1, :, 0:max_len, 0:max_len] # layers, batch, heads, seq_len, seq_len\n",
    "    tensor_decoder_attn = tensor_decoder_attn[:, batch:batch+1, :, 0:max_len, 0:max_len] # layers, batch, heads, seq_len, seq_len\n",
    "\n",
    "    # Return the formatted tokens and attention tensors\n",
    "    return tokens_input, tokens_output, tensor_encoder_attn, tensor_cross_attn, tensor_decoder_attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "cdf4df39-943a-446f-819a-e97e93129407",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens_input,\\\n",
    "tokens_output,\\\n",
    "tensor_encoder_attn,\\\n",
    "tensor_cross_attn,\\\n",
    "tensor_decoder_attn = format_attn(input, outputs, attentions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "66cbf390-5e54-4fd6-9274-e3c9d19d4ce1",
   "metadata": {},
   "outputs": [],
   "source": [
    "html_model_view = model_view(\n",
    "    encoder_attention=tensor_encoder_attn,\n",
    "    decoder_attention=tensor_decoder_attn,\n",
    "    cross_attention=tensor_cross_attn,\n",
    "    encoder_tokens=tokens_input[0:tensor_decoder_attn.size(-1)],\n",
    "    decoder_tokens=tokens_output[0:tensor_decoder_attn.size(-1)],\n",
    "    html_action='return'\n",
    ")\n",
    "with open(\"./out/model_view.html\", 'w') as file:\n",
    "    file.write(html_model_view.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "58589e22-44ec-4cd1-bf40-42cc8ed65764",
   "metadata": {},
   "outputs": [],
   "source": [
    "html_head_view = head_view(\n",
    "    encoder_attention=tensor_encoder_attn,\n",
    "    decoder_attention=tensor_decoder_attn,\n",
    "    cross_attention=tensor_cross_attn,\n",
    "    encoder_tokens=tokens_input[0:tensor_decoder_attn.size(-1)],\n",
    "    decoder_tokens=tokens_output[0:tensor_decoder_attn.size(-1)],\n",
    "    html_action='return'\n",
    ")\n",
    "with open(\"./out/head_view.html\", 'w') as file:\n",
    "    file.write(html_head_view.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "224bf51c-9872-415b-8d4f-1d6130bfe1d1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2016b47-3c21-4c23-af1e-4d7c39ba273e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dc936ea-0246-4fdc-b1d7-817a4c5f840e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}