{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sessions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NDIF uses a queue to handle concurrent requests from multiple users. To optimize the execution of our experiments we can use the `session` context to efficiently package multiple interventions together as one single request to the server.\n",
    "\n",
    "This offers the following benefits: \n",
    "1) All interventions within a session will be executed one after another without additional wait in the queue\n",
    "2) All intermediate outputs of each intervention are stored on the server and can be accessed by other interventions in the same session without moving the data back and forth between NDIF and the local machine.\n",
    "\n",
    "Let's take a look:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/anaconda3/envs/nnsight_local/lib/python3.12/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"
     ]
    }
   ],
   "source": [
    "from nnsight import LanguageModel, CONFIG\n",
    "import os\n",
    "\n",
    "# we are using Llama model remotely hosted on NDIF servers\n",
    "CONFIG.set_default_api_key(\"YOUR_API_KEY\") \n",
    "os.environ['HF_TOKEN'] = \"YOUR_HUGGING_FACE_TOKEN\"\n",
    "\n",
    "model = LanguageModel(\"meta-llama/Meta-Llama-3.1-70B\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-08-30 08:15:49,072 MainProcess nnsight_remote INFO     1478bff8-622a-42cf-9ba0-78d3a0803a8b - RECEIVED: Your job has been received and is waiting approval.\n",
      "2024-08-30 08:15:49,097 MainProcess nnsight_remote INFO     1478bff8-622a-42cf-9ba0-78d3a0803a8b - APPROVED: Your job was approved and is waiting to be run.\n",
      "2024-08-30 08:15:49,114 MainProcess nnsight_remote INFO     1478bff8-622a-42cf-9ba0-78d3a0803a8b - RUNNING: Your job has started running.\n",
      "2024-08-30 08:15:49,423 MainProcess nnsight_remote INFO     1478bff8-622a-42cf-9ba0-78d3a0803a8b - COMPLETED: Your job has been completed.\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "Downloading result: 100%|██████████| 1.69k/1.69k [00:00<00:00, 2.20MB/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "T1 - Original Prediction:   Paris\n",
      "T2 - Modified Prediction:   Paris\n"
     ]
    }
   ],
   "source": [
    "with model.session(remote=True) as session:\n",
    "\n",
    "  with model.trace(\"The Eiffel Tower is in the city of\") as t1:\n",
    "    # capture the hidden state from layer 11 at the last token\n",
    "    hs_79 = model.model.layers[79].output[0][:, -1, :] # no .save()\n",
    "    t1_tokens_out = model.lm_head.output.argmax(dim=-1).save()\n",
    "\n",
    "  with model.trace(\"Buckingham Palace is in the city of\") as t2:\n",
    "    model.model.layers[1].output[0][:, -1, :] = hs_79[:]\n",
    "    t2_tokens_out = model.lm_head.output.argmax(dim=-1).save()\n",
    "\n",
    "print(\"\\nT1 - Original Prediction: \", model.tokenizer.decode(t1_tokens_out[0][-1]))\n",
    "print(\"T2 - Modified Prediction: \", model.tokenizer.decode(t2_tokens_out[0][-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the example above, we are interested in replacing the hidden state of a later layer with an earlier one. Since we are using a `session`, we don't have to save the hidden state from Tracer 1 to reference it in Tracer 2.\n",
    "\n",
    "It is important to note that all the traces defined within the `session` context are executed sequentially, strictly following the order of definition (i.e. `t2` being executed after `t1` and `t3` after `t2` etc.)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `session` context object has its own methods to log values and be terminated early."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-08-30 08:15:53,770 MainProcess nnsight_remote INFO     e794cd1a-3c4d-4300-ba68-6dcb6d8bda02 - RECEIVED: Your job has been received and is waiting approval.\n",
      "2024-08-30 08:15:53,795 MainProcess nnsight_remote INFO     e794cd1a-3c4d-4300-ba68-6dcb6d8bda02 - APPROVED: Your job was approved and is waiting to be run.\n",
      "2024-08-30 08:15:53,801 MainProcess nnsight_remote INFO     e794cd1a-3c4d-4300-ba68-6dcb6d8bda02 - RUNNING: Your job has started running.\n",
      "2024-08-30 08:15:53,811 MainProcess nnsight_remote INFO     e794cd1a-3c4d-4300-ba68-6dcb6d8bda02 - LOG: -- Early Stop --\n",
      "2024-08-30 08:15:53,823 MainProcess nnsight_remote INFO     e794cd1a-3c4d-4300-ba68-6dcb6d8bda02 - COMPLETED: Your job has been completed.\n",
      "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 5.66MB/s]\n"
     ]
    }
   ],
   "source": [
    "with model.session(remote=True) as session:\n",
    "\n",
    "  session.log(\"-- Early Stop --\")\n",
    "  session.exit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to the benefits mentioned above, the `session` context also enables interesting experiments not possible with other `nnsight` tools - since every trace is run on its own model, it means that within one session we can run interventions between different models – for example, we can swap activations between vanilla and instruct versions of the Llama model and compare the outputs. And `session` can also be used to run experiments entirely locally! "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nnsight_local",
   "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
