{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3d129fe7-5a5f-441f-bb3f-288073f30c44",
   "metadata": {},
   "source": [
    "# Multiple Token Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39793d68-c96e-4069-82a3-31c283eb08eb",
   "metadata": {},
   "source": [
    "When generating more than one token, use `<module>.next()` to denote following interventions should be applied to the subsequent generations for that module.\n",
    "\n",
    "Here we generate three tokens and save the hidden states of the last layer for each one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "032125c8-3d90-4671-bd66-4cbbf090c871",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nnsight import LanguageModel\n",
    "\n",
    "model = LanguageModel('openai-community/gpt2', device_map='cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a48b0c3a-ac2d-4e5c-9dbf-cf1df44de0ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
      "A decoder-only architecture is being used, but right-padding was detected! For correct generation results, please set `padding_side='left'` when initializing the tokenizer.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    }
   ],
   "source": [
    "with model.generate('The Eiffel Tower is in the city of', max_new_tokens=3) as tracer:\n",
    "\n",
    "    hidden_states1 = model.transformer.h[-1].output[0].save()\n",
    "    \n",
    "    hidden_states2 = model.transformer.h[-1].next().output[0].save()\n",
    "    \n",
    "    hidden_states3 = model.transformer.h[-1].next().output[0].save()\n",
    "\n",
    "    out = model.generator.output.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0df3f8a-b8fd-4b9d-9f99-3bef39a7143f",
   "metadata": {},
   "source": [
    "Note how calling save *before* `tracer.next()` returns the hidden state across the initial prompt while calling save *after* returns the hidden state of each subsequent generated token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cea3a24a-773d-4305-a04d-2e964680c381",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 10, 768])\n",
      "torch.Size([1, 1, 768])\n",
      "torch.Size([1, 1, 768])\n",
      "tensor([[ 464,  412,  733,  417, 8765,  318,  287,  262, 1748,  286, 6342,   11,\n",
      "          290]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "print(hidden_states1.shape)\n",
    "print(hidden_states2.shape)\n",
    "print(hidden_states3.shape)\n",
    "print(out)"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
