{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5873a43e-e181-47d3-8ba9-2f6c9e7a7d61",
   "metadata": {},
   "source": [
    "# Cross-Prompt Intervention"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fb6417c-53dd-43d4-900b-7fdb9616591c",
   "metadata": {},
   "source": [
    "Intervention operations work cross prompt! Use two invocations within the same generation block and operations can work between them.\n",
    "\n",
    "In this case, we grab the token embeddings coming from the first prompt, \"Madison square garden is located in the city of New\" and replace the embeddings of the second prompt with them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "455cfa6f-54d6-43b7-befd-6754d6baa9e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/rockinglikebalboa/NDIF/nnsight-test-env/lib/python3.10/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",
      "/Users/rockinglikebalboa/NDIF/nnsight-test-env/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from nnsight import LanguageModel\n",
    "\n",
    "model = LanguageModel('openai-community/gpt2', device_map='auto')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "50c3b04f-a9dc-4f01-9010-a93751745a47",
   "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",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Madison square garden is located in the city of New York City.']\n",
      "['_ _ _ _ _ _ _ _ _ _ York City.']\n"
     ]
    }
   ],
   "source": [
    "with model.generate(max_new_tokens=3) as tracer:\n",
    "    \n",
    "    with tracer.invoke(\"Madison square garden is located in the city of New\") as invoker:\n",
    "\n",
    "        embeddings = model.transformer.wte.output\n",
    "        original = model.generator.output.save()\n",
    "\n",
    "    with tracer.invoke(\"_ _ _ _ _ _ _ _ _ _\") as invoker:\n",
    "\n",
    "        model.transformer.wte.output = embeddings\n",
    "        intervened = model.generator.output.save()\n",
    "\n",
    "print(model.tokenizer.batch_decode(original.value))\n",
    "print(model.tokenizer.batch_decode(intervened.value))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c70184c-239d-4964-94bc-56f5d4fea363",
   "metadata": {},
   "source": [
    "We also could have entered a pre-saved embedding tensor as shown here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "808f77a3-a43c-4564-a47b-0f1b1da6c79b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Madison square garden is located in the city of New York City.']\n",
      "['_ _ _ _ _ _ _ _ _ _ York City.']\n"
     ]
    }
   ],
   "source": [
    "with model.generate(\"Madison square garden is located in the city of New\", max_new_tokens=3) as tracer:\n",
    "\n",
    "    embeddings = model.transformer.wte.output.save()\n",
    "    original = model.generator.output.save()\n",
    "\n",
    "print(model.tokenizer.batch_decode(original.value))\n",
    "\n",
    "with model.generate(\"_ _ _ _ _ _ _ _ _ _\", max_new_tokens=3) as tracer:\n",
    "\n",
    "    model.transformer.wte.output = embeddings.value\n",
    "    intervened = model.generator.output.save()\n",
    "\n",
    "print(model.tokenizer.batch_decode(intervened.value))"
   ]
  }
 ],
 "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
}
