{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Activation Oracle via the vllm-lens Inspect Provider\n",
    "\n",
    "This notebook demonstrates the **activation oracle** using the `vllm-lens`\n",
    "[Inspect AI](https://inspect.aisi.org.uk/) model provider, structured as an\n",
    "Inspect **solver** so it can be composed into evaluation tasks.\n",
    "\n",
    "The oracle (from [LatentQA / Activation Oracles](https://arxiv.org/abs/2412.08586))\n",
    "lets you \"ask questions\" about a model's internal activations.  A LoRA adapter\n",
    "interprets steering vectors injected at a special token, enabling the model to\n",
    "describe what it is \"thinking about\" at a given layer and token.\n",
    "\n",
    "**Flow (inside the solver):**\n",
    "1. Capture residual-stream activations from a target prompt.\n",
    "2. Extract the activation at a token of interest (e.g. \"pupil\").\n",
    "3. Inject it into an oracle prompt via steering and generate with the oracle LoRA."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from inspect_ai import Task, eval\n",
    "from inspect_ai.dataset import MemoryDataset, Sample\n",
    "from inspect_ai.model import (\n",
    "    ChatMessageAssistant,\n",
    "    ChatMessageUser,\n",
    "    GenerateConfig,\n",
    "    ChatMessage,\n",
    "    get_model,\n",
    ")\n",
    "from inspect_ai.solver import Generate, TaskState, solver\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "from vllm_lens import SteeringVector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL = \"Qwen/Qwen3-8B\"\n",
    "LORA_PATH = \"adamkarvonen/checkpoints_latentqa_cls_past_lens_addition_Qwen3-8B\"\n",
    "TARGET_CONTENT = (\n",
    "    \"The philosopher who drank hemlock taught a student who founded \"\n",
    "    \"an academy. That student's most famous pupil was\"\n",
    ")\n",
    "ORACLE_QUESTION = \"Can you name all people the model is currently thinking about?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start the Model and Capture Activations\n",
    "\n",
    "Start the vLLM server once with LoRA support enabled, then run a quick\n",
    "activation capture to verify the residual stream extraction works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "residual_stream shape: torch.Size([3, 24, 4096])\n"
     ]
    }
   ],
   "source": [
    "model = get_model(\n",
    "    f\"vllm-lens/{MODEL}\",\n",
    "    enable_lora=True,\n",
    "    max_lora_rank=64,\n",
    "    max_model_len=4096,\n",
    "    gpu_memory_utilization=0.85,\n",
    "    lora_modules=f\"oracle={LORA_PATH}\",\n",
    ")\n",
    "capture_config = GenerateConfig(\n",
    "    temperature=0.0,\n",
    "    max_tokens=1,\n",
    "    extra_body={\n",
    "        \"extra_args\": {\"output_residual_stream\": [9, 18, 27]},\n",
    "        \"chat_template_kwargs\": {\"enable_thinking\": False},\n",
    "    },\n",
    ")\n",
    "output = await model.generate(\n",
    "    [ChatMessageAssistant(content=TARGET_CONTENT)], config=capture_config\n",
    ")\n",
    "residual_stream = output.metadata[\"activations\"][\"residual_stream\"]  # type: ignore\n",
    "print(\"residual_stream shape:\", residual_stream.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the Activation Oracle Solver\n",
    "\n",
    "The solver extracts the target text from the sample's assistant message,\n",
    "captures residual-stream activations, locates the token of interest, and\n",
    "then queries the oracle LoRA at each captured layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@solver\n",
    "def ask_activation_oracle(\n",
    "    oracle_question: str,\n",
    "    token_of_interest: str,\n",
    "    lora_path: str,\n",
    "    extraction_activation_layers: list[int] = [9, 18, 27],\n",
    "    injection_layer: int = 1,\n",
    "    injection_special_token: str = \" ?\",\n",
    "    steering_coefficient: float = 1.0,\n",
    "):\n",
    "    tokenizer = AutoTokenizer.from_pretrained(MODEL)\n",
    "\n",
    "    async def solve(state: TaskState, _solver_generate: Generate) -> TaskState:\n",
    "        # Extract activations from the forward pass\n",
    "        print(state.messages)\n",
    "        model = get_model()\n",
    "        capture_config = GenerateConfig(\n",
    "            temperature=0.0,\n",
    "            max_tokens=1,\n",
    "            extra_body={\n",
    "                \"extra_args\": {\"output_residual_stream\": extraction_activation_layers},\n",
    "                \"chat_template_kwargs\": {\"enable_thinking\": False},\n",
    "            },\n",
    "        )\n",
    "        output = await model.generate(state.messages, config=capture_config)\n",
    "        residual_stream = output.metadata[\"activations\"][\"residual_stream\"]\n",
    "        print(\"residual_stream shape:\", residual_stream.shape)\n",
    "\n",
    "        # --- Step 2: Find token of interest and extract activations ---\n",
    "        target_prompt = tokenizer.apply_chat_template(\n",
    "            state.messages,\n",
    "            tokenize=False,\n",
    "            add_generation_prompt=True,\n",
    "            enable_thinking=False,\n",
    "        )\n",
    "        target_token_ids = tokenizer.encode(target_prompt, add_special_tokens=False)\n",
    "\n",
    "        tok_candidates = tokenizer.encode(\n",
    "            f\" {token_of_interest}\", add_special_tokens=False\n",
    "        )\n",
    "        tok_id = tok_candidates[-1]\n",
    "\n",
    "        tok_pos = None\n",
    "        for i, tid in enumerate(target_token_ids):\n",
    "            if tid == tok_id:\n",
    "                tok_pos = i\n",
    "        assert tok_pos is not None, f\"Could not find '{token_of_interest}' token\"\n",
    "\n",
    "        token_acts = {}\n",
    "        for layer_i, layer in enumerate(extraction_activation_layers):\n",
    "            token_acts[layer] = residual_stream[layer_i, tok_pos, :]\n",
    "\n",
    "        # --- Step 3: Run oracle at each layer ---\n",
    "        oracle_results = []\n",
    "        for layer in extraction_activation_layers:\n",
    "            prefix = f\"Layer: {layer}\\n{injection_special_token} \\n\"\n",
    "            oracle_content = prefix + oracle_question\n",
    "\n",
    "            oracle_messages = [{\"role\": \"user\", \"content\": oracle_content}]\n",
    "            oracle_prompt = tokenizer.apply_chat_template(\n",
    "                oracle_messages,\n",
    "                tokenize=False,\n",
    "                add_generation_prompt=True,\n",
    "                enable_thinking=False,\n",
    "            )\n",
    "            oracle_token_ids = tokenizer.encode(oracle_prompt, add_special_tokens=False)\n",
    "\n",
    "            special_token_id = tokenizer.encode(\n",
    "                injection_special_token, add_special_tokens=False\n",
    "            )\n",
    "            assert len(special_token_id) == 1\n",
    "            special_pos = oracle_token_ids.index(special_token_id[0])\n",
    "\n",
    "            act_vector = token_acts[layer].unsqueeze(0).unsqueeze(0)\n",
    "\n",
    "            messages: list[ChatMessage] = [ChatMessageUser(content=oracle_content)]\n",
    "\n",
    "            oracle_config = GenerateConfig(\n",
    "                temperature=0.0,\n",
    "                max_tokens=50,\n",
    "                extra_body={\n",
    "                    \"extra_args\": {\n",
    "                        \"apply_steering_vectors\": [\n",
    "                            SteeringVector(\n",
    "                                activations=act_vector,\n",
    "                                layer_indices=[injection_layer],\n",
    "                                scale=steering_coefficient,\n",
    "                                norm_match=True,\n",
    "                                position_indices=[special_pos],\n",
    "                            )\n",
    "                        ],\n",
    "                    },\n",
    "                    \"lora_request\": {\n",
    "                        \"lora_name\": \"oracle\",\n",
    "                        \"lora_int_id\": 1,\n",
    "                        \"lora_path\": lora_path,\n",
    "                    },\n",
    "                    \"chat_template_kwargs\": {\"enable_thinking\": False},\n",
    "                },\n",
    "            )\n",
    "\n",
    "            response = await model.generate(messages, config=oracle_config)\n",
    "\n",
    "            oracle_results.append(\n",
    "                {\n",
    "                    \"layer\": layer,\n",
    "                    \"activation_norm\": float(token_acts[layer].norm().item()),\n",
    "                    \"response\": response.completion,\n",
    "                }\n",
    "            )\n",
    "\n",
    "        state.metadata[\"oracle_results\"] = oracle_results\n",
    "        return state\n",
    "\n",
    "    return solve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Up Task and Run\n",
    "\n",
    "Create a `Task` with a single `Sample` whose input is the target content\n",
    "(as an assistant message), wire up the solver, and run with `eval()`.\n",
    "The model started above is reused for the evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f5f0cedcca594856bba571c76064e50e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sample = Sample(\n",
    "    input=[ChatMessageAssistant(content=TARGET_CONTENT)],\n",
    "    target=\"oracle\",\n",
    ")\n",
    "\n",
    "task = Task(\n",
    "    dataset=MemoryDataset([sample]),\n",
    "    solver=ask_activation_oracle(\n",
    "        oracle_question=ORACLE_QUESTION,\n",
    "        token_of_interest=\"pupil\",\n",
    "        lora_path=LORA_PATH,\n",
    "    ),  # type: ignore\n",
    ")\n",
    "\n",
    "log = eval(task, model=model, display=\"none\")[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Status: success\n",
      "\n",
      "=== Layer 9 (activation norm: 69.50) ===\n",
      "Oracle: The assistant is currently thinking about the philosopher NAME_1.\n",
      "\n",
      "=== Layer 18 (activation norm: 110.50) ===\n",
      "Oracle: The model is currently thinking about Socrates, Plato, and Aristotle.\n",
      "\n",
      "=== Layer 27 (activation norm: 368.00) ===\n",
      "Oracle: The assistant is currently thinking about the philosopher NAME_1.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(f\"Status: {log.status}\\n\")\n",
    "\n",
    "for result in log.samples[0].metadata[\"oracle_results\"]:  # type: ignore\n",
    "    layer = result[\"layer\"]\n",
    "    norm = result[\"activation_norm\"]\n",
    "    response = result[\"response\"]\n",
    "    print(f\"=== Layer {layer} (activation norm: {norm:.2f}) ===\")\n",
    "    print(f\"Oracle: {response}\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vllm-lens (3.12.13)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}