{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setting Values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We often not only want to see whats happening during computation, but intervene and edit the flow of information.\n",
    "\n",
    "In this example, we create a tensor of noise to add to the hidden states. We then add it, use the assigment `=` operator to update the tensors of `.output[0][:]` with these new noised values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/share/u/caden/.conda/envs/interp/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "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_pre = model.transformer.h[-1].output[0].clone().save()\n",
    "\n",
    "    noise = (0.001**0.5)*torch.randn(hidden_states_pre.shape)\n",
    "\n",
    "    # model.transformer.h[-1].output = (hidden_states_pre + noise, model.transformer.h[-1].output[1])\n",
    "    model.transformer.h[-1].output[0][:] = hidden_states_pre + noise\n",
    "\n",
    "    hidden_states_post = model.transformer.h[-1].output[0].save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the change in the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.0505, -0.1728, -0.1690,  ..., -1.0096,  0.1280, -1.0687],\n",
      "         [ 8.7494,  2.9057,  5.3024,  ..., -8.0418,  1.2964, -2.8677],\n",
      "         [ 0.2960,  4.6686, -3.6642,  ...,  0.2391, -2.6064,  3.2263],\n",
      "         ...,\n",
      "         [ 2.1537,  6.8917,  3.8651,  ...,  0.0588, -1.9866,  5.9188],\n",
      "         [-0.4460,  7.4285, -9.3065,  ...,  2.0528, -2.7946,  0.5556],\n",
      "         [ 6.6286,  1.7258,  4.7969,  ...,  7.6714,  3.0683,  2.0481]]],\n",
      "       device='cuda:0', grad_fn=<CloneBackward0>)\n",
      "tensor([[[-0.0244, -0.1430, -0.1384,  ..., -1.0185,  0.1173, -1.0855],\n",
      "         [ 8.8112,  2.8995,  5.3000,  ..., -8.0057,  1.3378, -2.8778],\n",
      "         [ 0.2780,  4.6427, -3.6233,  ...,  0.1378, -2.5877,  3.2030],\n",
      "         ...,\n",
      "         [ 2.1590,  6.9278,  3.8584,  ...,  0.0728, -2.0145,  5.9685],\n",
      "         [-0.4962,  7.4435, -9.3371,  ...,  2.1299, -2.7784,  0.5379],\n",
      "         [ 6.6941,  1.7279,  4.8105,  ...,  7.6403,  3.0599,  2.0637]]],\n",
      "       device='cuda:0', grad_fn=<CopySlices>)\n"
     ]
    }
   ],
   "source": [
    "print(hidden_states_pre)\n",
    "print(hidden_states_post)"
   ]
  }
 ],
 "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
}
