{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another thing we can do is apply modules in the model's module tree at any point during computation, even if it's out of order.\n",
    "\n",
    "Here we get the hidden states of the last layer like usual. We also chain apply `model.transformer.ln_f` and `model.lm_head` in order to \"decode\" the hidden states into vocabularly space. Applying softmax and then argmax allows us to then transform the vocabulary space hidden states into actually tokens which we can then use the tokenizer to decode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a GPT2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    }
   ],
   "source": [
    "from nnsight import LanguageModel\n",
    "import torch\n",
    "\n",
    "model = LanguageModel(\"openai-community/gpt2\", device_map='cuda')\n",
    "\n",
    "with model.trace('The Eiffel Tower is in the city of') as tracer:\n",
    "        \n",
    "    hidden_states = model.transformer.h[-1].output[0]\n",
    "\n",
    "    hidden_states = model.lm_head(model.transformer.ln_f(hidden_states)).save()\n",
    "    \n",
    "    tokens = torch.softmax(hidden_states, dim=2).argmax(dim=2).save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ -36.2875,  -35.0115,  -38.0794,  ...,  -40.5164,  -41.3760,\n",
      "           -34.9194],\n",
      "         [ -68.8886,  -70.1562,  -71.8408,  ...,  -80.4194,  -78.2552,\n",
      "           -71.1206],\n",
      "         [ -82.2950,  -81.6519,  -83.9940,  ...,  -94.4878,  -94.5194,\n",
      "           -85.6997],\n",
      "         ...,\n",
      "         [-113.8675, -111.8628, -113.6634,  ..., -116.7652, -114.8267,\n",
      "          -112.3621],\n",
      "         [ -81.8531,  -83.3007,  -91.8193,  ...,  -92.9943,  -89.8382,\n",
      "           -85.6898],\n",
      "         [-103.9307, -102.5054, -105.1563,  ..., -109.3099, -110.4196,\n",
      "          -103.1395]]], device='cuda:0', grad_fn=<UnsafeViewBackward0>)\n",
      "tensor([[ 198,   12,  417, 8765,  318,  257,  262, 3504, 7372, 6342]],\n",
      "       device='cuda:0')\n",
      "\n",
      "-el Tower is a the middle centre Paris\n"
     ]
    }
   ],
   "source": [
    "print(hidden_states)\n",
    "print(tokens)\n",
    "print(model.tokenizer.decode(tokens[0]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ndif",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
