{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook 1: Find the Direction\n",
    "\n",
    "**AI Alignment — The Multiverse School**\n",
    "\n",
    "Every concept a language model has learned lives somewhere in its activation space as a *direction* — a vector.\n",
    "Honesty, helpfulness, harm avoidance — they're not single neurons, they're geometric features spread across thousands of dimensions.\n",
    "\n",
    "Today you will **find one of those directions**, extract it, and prove it's real.\n",
    "\n",
    "**What you need:** A vast.ai GPU instance (setup instructions below)"
   ],
   "id": "cell-00"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0 — vast.ai Setup\n",
    "\n",
    "You're running this on your own rented GPU, not Colab. Here's how to get there.\n",
    "\n",
    "### Rent an instance\n",
    "\n",
    "1. Go to [cloud.vast.ai](https://cloud.vast.ai) and sign up (or log in)\n",
    "2. Add credits ($5-10 is plenty for today)\n",
    "3. Under **Account → Keys**, add your SSH public key (the *contents* of `~/.ssh/id_ed25519.pub`, starting with `ssh-ed25519`)\n",
    "4. Go to **Templates** and pick **PyTorch (Vast)** — it comes with torch, CUDA, and Jupyter\n",
    "5. Set disk space to **40 GB** (model weights are ~9 GB)\n",
    "6. Go to **Search** and pick a GPU with at least 16 GB VRAM. Any of these work for Qwen3.5-4B:\n",
    "   - **RTX 3090** (24 GB, ~$0.15-0.30/hr) — cheapest option, recommended\n",
    "   - **RTX 4090** (24 GB, ~$0.30-0.50/hr) — faster\n",
    "   - **A40 / A6000** (48 GB) — overkill but fine\n",
    "7. Filter for **reliability > 98%** and click **Rent**\n",
    "\n",
    "### Open Jupyter\n",
    "\n",
    "Once your instance is running:\n",
    "1. Go to **Instances** and click **Open** (Jupyter)\n",
    "2. Upload this notebook (drag it into the file browser)\n",
    "3. Open it and start running cells\n",
    "\n",
    "### Or use the terminal / your own Jupyter\n",
    "\n",
    "The **Connect** button on the instance shows an SSH command like:\n",
    "```\n",
    "ssh -p <port> root@<instance-ip> -L 8080:localhost:8080\n",
    "```\n",
    "The `-L` flag tunnels the instance's Jupyter to `http://localhost:8080` on your laptop."
   ],
   "id": "cell-01"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Check your GPU — you should see your vast.ai GPU here\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "import torch\n",
    "\n",
    "print(f\"PyTorch: {torch.__version__}\")\n",
    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n",
    "else:\n",
    "    print(\"WARNING: No GPU detected. Check your vast.ai instance.\")\n",
    "\n",
    "# Install what's not in the PyTorch template (Qwen3.5 needs transformers 5.x)\n",
    "subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"-U\",\n",
    "                       \"transformers>=5.17\", \"accelerate\", \"scikit-learn\", \"matplotlib\"])\n",
    "print(\"\\nDependencies installed. If transformers was upgraded, restart the kernel before continuing.\")"
   ],
   "id": "cell-02"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# --- Langfuse tracing (optional) ---\n",
    "# Put LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY and LANGFUSE_BASE_URL in a .env file next to this notebook.\n",
    "# Optionally set STUDENT_ID in .env so your runs are grouped under your name.\n",
    "# With no keys, everything below still runs; it just isn't traced.\n",
    "subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"langfuse>=4\", \"python-dotenv\"])\n",
    "\n",
    "import os\n",
    "from contextlib import contextmanager\n",
    "from datetime import datetime\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "NOTEBOOK = \"01_probe_find_the_direction\"\n",
    "SESSION_ID = f\"{NOTEBOOK}-{datetime.now():%Y%m%d-%H%M%S}\"\n",
    "\n",
    "class _NoTrace:\n",
    "    def update(self, **kwargs): pass\n",
    "    def score(self, **kwargs): pass\n",
    "\n",
    "TRACING = False\n",
    "if os.environ.get(\"LANGFUSE_SECRET_KEY\"):\n",
    "    from langfuse import get_client, propagate_attributes\n",
    "    langfuse = get_client()\n",
    "    TRACING = langfuse.auth_check()\n",
    "\n",
    "@contextmanager\n",
    "def trace(name, as_type=\"span\", **kwargs):\n",
    "    \"\"\"A Langfuse observation (span or generation) grouped under this notebook run. No-op without keys.\"\"\"\n",
    "    if not TRACING:\n",
    "        yield _NoTrace()\n",
    "        return\n",
    "    with propagate_attributes(session_id=SESSION_ID, user_id=os.environ.get(\"STUDENT_ID\"),\n",
    "                              tags=[\"ai-alignment\", NOTEBOOK]):\n",
    "        with langfuse.start_as_current_observation(name=name, as_type=as_type, **kwargs) as obs:\n",
    "            yield obs\n",
    "\n",
    "print(f\"Langfuse tracing: {'on, session ' + SESSION_ID if TRACING else 'off (no keys in .env)'}\")"
   ],
   "id": "cell-03"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 — Load the Model\n",
    "\n",
    "We're using **Qwen3.5-4B** — a 4-billion-parameter instruction-tuned model from Alibaba's Qwen team,\n",
    "the smallest-but-still-capable member of the newest Qwen family.\n",
    "\n",
    "In bfloat16 it's about 9 GB, which fits on any GPU you picked above.\n",
    "\n",
    "When the model loads, look at the numbers it prints. You're about to reach inside\n",
    "those layers and find what the model uses to represent honesty."
   ],
   "id": "cell-04"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "MODEL_NAME = \"Qwen/Qwen3.5-4B\"\n",
    "\n",
    "print(\"Loading tokenizer...\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "\n",
    "print(\"Loading model (the first load downloads ~9 GB; it's cached after that)...\")\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    dtype=torch.bfloat16,\n",
    "    device_map=\"auto\",\n",
    ")\n",
    "model.eval()\n",
    "\n",
    "print(f\"\\n--- Model loaded ---\")\n",
    "print(f\"Parameters: {sum(p.numel() for p in model.parameters()) / 1e9:.1f} billion\")\n",
    "print(f\"Layers: {model.config.num_hidden_layers}\")\n",
    "print(f\"Hidden size: {model.config.hidden_size}\")\n",
    "print(f\"Device: {model.device}\")\n",
    "print(f\"GPU memory used: {torch.cuda.memory_allocated() / 1e9:.1f} GB\")\n",
    "print(f\"\\nEvery layer produces a {model.config.hidden_size}-dimensional vector.\")\n",
    "print(f\"The concept of 'honesty' is a direction in that {model.config.hidden_size}-D space.\")\n",
    "print(f\"We're going to find it.\")"
   ],
   "id": "cell-05"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 — What's Inside a Layer?\n",
    "\n",
    "Before we go looking for honesty, let's see what a hidden state actually looks like.\n",
    "We'll feed the model a simple prompt and look at the raw activations."
   ],
   "id": "cell-06"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Feed the model a prompt and capture ALL hidden states from ALL layers\n",
    "test_input = \"Is this claim actually true?\"\n",
    "inputs = tokenizer(test_input, return_tensors=\"pt\").to(model.device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    outputs = model(**inputs, output_hidden_states=True)\n",
    "\n",
    "# outputs.hidden_states is a tuple: the embeddings, then one tensor per layer.\n",
    "# So hidden_states[i + 1] is the OUTPUT of layer i.\n",
    "print(f\"Number of hidden states: {len(outputs.hidden_states)} (embeddings + {model.config.num_hidden_layers} layers)\")\n",
    "print(f\"Shape of each: {outputs.hidden_states[0].shape}\")\n",
    "print(f\"  → [batch_size=1, sequence_length={inputs['input_ids'].shape[1]}, hidden_dim={model.config.hidden_size}]\")\n",
    "\n",
    "mid_layer = model.config.num_hidden_layers // 2\n",
    "print(f\"\\nThe last token's hidden state coming out of layer {mid_layer}:\")\n",
    "last_token_state = outputs.hidden_states[mid_layer + 1][0, -1, :]\n",
    "print(f\"  Shape: {last_token_state.shape}\")\n",
    "print(f\"  First 10 values: {last_token_state[:10].float().cpu().tolist()}\")\n",
    "print(f\"\\nThat vector IS what the model is 'thinking' at this point.\")\n",
    "print(f\"Somewhere in those {model.config.hidden_size} dimensions is the direction for honesty.\")"
   ],
   "id": "cell-07"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 — Contrastive Pairs\n",
    "\n",
    "To find the honesty direction, we need to show the model the **same scenario**\n",
    "framed two ways: one that activates honesty, one that activates sycophancy\n",
    "(telling people what they want to hear).\n",
    "\n",
    "The difference in hidden states between these two framings — averaged across\n",
    "many examples — IS the honesty direction.\n",
    "\n",
    "**Read through these pairs.** Each one is the same question, but the system\n",
    "prompt tells the model to be either honest or people-pleasing."
   ],
   "id": "cell-08"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# 20 matched pairs: same user question, different system framing\n",
    "# The model's hidden states will differ — and that difference is what we're after.\n",
    "\n",
    "SCENARIOS = [\n",
    "    \"Is this too-good-to-be-true investment actually safe?\",\n",
    "    \"What are the real side effects of this medication?\",\n",
    "    \"Is this online review trustworthy?\",\n",
    "    \"Does this weight loss supplement really work?\",\n",
    "    \"Is my data actually private on this platform?\",\n",
    "    \"Will this coding bootcamp really get me a job?\",\n",
    "    \"Is this charity using donations effectively?\",\n",
    "    \"Are these lab results actually concerning?\",\n",
    "    \"Is this contractor giving me a fair quote?\",\n",
    "    \"Does this academic paper actually prove its claim?\",\n",
    "    \"Is this news article reporting facts accurately?\",\n",
    "    \"Will this skincare product actually reduce wrinkles?\",\n",
    "    \"Is this used car in good condition?\",\n",
    "    \"Is my password really secure enough?\",\n",
    "    \"Does this insurance policy actually cover what I need?\",\n",
    "    \"Is this organic label meaningful or just marketing?\",\n",
    "    \"Will this online course teach me real skills?\",\n",
    "    \"Is this financial advisor acting in my interest?\",\n",
    "    \"Are these product reviews genuine?\",\n",
    "    \"Is this job offer legitimate?\",\n",
    "]\n",
    "\n",
    "HONEST_SYS = \"You are an honest, straightforward assistant.\"\n",
    "SYCOPHANTIC_SYS = \"You are a people-pleasing assistant who always tells users what they want to hear.\"\n",
    "\n",
    "honest_prompts = [f\"{HONEST_SYS} User asks: {s}\" for s in SCENARIOS]\n",
    "sycophantic_prompts = [f\"{SYCOPHANTIC_SYS} User asks: {s}\" for s in SCENARIOS]\n",
    "\n",
    "print(f\"Created {len(honest_prompts)} honest and {len(sycophantic_prompts)} sycophantic prompts.\")\n",
    "print(f\"\\nExample honest:      {honest_prompts[0][:80]}...\")\n",
    "print(f\"Example sycophantic: {sycophantic_prompts[0][:80]}...\")"
   ],
   "id": "cell-09"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4 — Extract Hidden States\n",
    "\n",
    "Now we feed every prompt through the model and capture the hidden state\n",
    "coming out of the **middle layer**, at the **last token position**.\n",
    "\n",
    "Why middle layer? Early layers handle syntax. Late layers handle output\n",
    "formatting. The middle layers are where *concepts* live — honesty, deception,\n",
    "helpfulness — as geometric features."
   ],
   "id": "cell-10"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def get_hidden_states(prompts, model, tokenizer, layer_idx):\n",
    "    \"\"\"Extract the output of layer `layer_idx` at the last token for each prompt.\"\"\"\n",
    "    states = []\n",
    "    for i, prompt in enumerate(prompts):\n",
    "        inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
    "        with torch.no_grad():\n",
    "            out = model(**inputs, output_hidden_states=True)\n",
    "        # +1 because hidden_states[0] is the embedding layer\n",
    "        hs = out.hidden_states[layer_idx + 1][:, -1, :].float().cpu().numpy()\n",
    "        states.append(hs[0])\n",
    "        if (i + 1) % 5 == 0:\n",
    "            print(f\"  {i+1}/{len(prompts)}\")\n",
    "    return np.array(states)\n",
    "\n",
    "mid_layer = model.config.num_hidden_layers // 2\n",
    "print(f\"Extracting from layer {mid_layer} (of {model.config.num_hidden_layers})...\")\n",
    "\n",
    "print(\"\\nHonest prompts:\")\n",
    "honest_states = get_hidden_states(honest_prompts, model, tokenizer, mid_layer)\n",
    "\n",
    "print(\"\\nSycophantic prompts:\")\n",
    "syco_states = get_hidden_states(sycophantic_prompts, model, tokenizer, mid_layer)\n",
    "\n",
    "print(f\"\\nHonest states shape:      {honest_states.shape}\")\n",
    "print(f\"Sycophantic states shape: {syco_states.shape}\")\n",
    "print(f\"\\nEach row is a {model.config.hidden_size}-dimensional snapshot of what the model\")\n",
    "print(f\"was 'thinking' after reading that prompt. 40 snapshots total: 20 honest, 20 sycophantic.\")"
   ],
   "id": "cell-11"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 — Train the Probe\n",
    "\n",
    "A **linear probe** is a simple classifier (logistic regression) trained on\n",
    "the hidden states. If it can tell honest from sycophantic states with high\n",
    "accuracy, that proves the model has a **linearly separable** internal\n",
    "representation of honesty.\n",
    "\n",
    "In plain terms: the model already knows the difference. We're just finding\n",
    "where it stores that knowledge."
   ],
   "id": "cell-12"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.model_selection import cross_val_score\n",
    "\n",
    "# Stack all states and label them\n",
    "X = np.concatenate([honest_states, syco_states])\n",
    "y = np.array([1] * len(honest_states) + [0] * len(syco_states))\n",
    "# 1 = honest, 0 = sycophantic\n",
    "\n",
    "# Train with 5-fold cross-validation\n",
    "probe = LogisticRegression(max_iter=1000)\n",
    "scores = cross_val_score(probe, X, y, cv=5)\n",
    "\n",
    "# Fit on all data (we'll use the weights as our direction vector)\n",
    "probe.fit(X, y)\n",
    "\n",
    "print(f\"Probe accuracy (5-fold CV): {scores.mean():.1%} ± {scores.std():.1%}\")\n",
    "with trace(\"train-honesty-probe\", input={\"layer\": mid_layer, \"n_examples\": len(X)},\n",
    "           metadata={\"model\": MODEL_NAME}) as obs:\n",
    "    obs.update(output={\"cv_accuracy_mean\": float(scores.mean()), \"cv_accuracy_std\": float(scores.std())})\n",
    "    obs.score(name=\"probe_cv_accuracy\", value=float(scores.mean()))\n",
    "print(f\"\\nWhat this means:\")\n",
    "if scores.mean() > 0.9:\n",
    "    print(f\"  The model's hidden states are {scores.mean():.0%} separable on honesty vs sycophancy.\")\n",
    "    print(f\"  The model HAS an internal representation of this concept.\")\n",
    "    print(f\"  It's not just a prompt trick — it's geometric.\")\n",
    "elif scores.mean() > 0.7:\n",
    "    print(f\"  Partially separable — the signal exists but it's noisy.\")\n",
    "else:\n",
    "    print(f\"  Weak separation — try a different layer or more contrastive pairs.\")"
   ],
   "id": "cell-13"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6 — Visualize the Separation\n",
    "\n",
    "Let's project those high-dimensional vectors down to 2D with PCA and actually\n",
    "*see* the honesty direction."
   ],
   "id": "cell-14"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "pca = PCA(n_components=2)\n",
    "X_2d = pca.fit_transform(X)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(8, 6))\n",
    "ax.scatter(X_2d[:len(honest_states), 0], X_2d[:len(honest_states), 1],\n",
    "           c='#2563a8', label='Honest', alpha=0.8, s=60, edgecolors='white', linewidth=0.5)\n",
    "ax.scatter(X_2d[len(honest_states):, 0], X_2d[len(honest_states):, 1],\n",
    "           c='#dc2626', label='Sycophantic', alpha=0.8, s=60, edgecolors='white', linewidth=0.5)\n",
    "\n",
    "# Draw the probe's decision boundary, projected into PCA space\n",
    "w = probe.coef_[0] @ pca.components_.T\n",
    "b = probe.intercept_[0] + probe.coef_[0] @ pca.mean_\n",
    "xx = np.linspace(X_2d[:, 0].min() - 1, X_2d[:, 0].max() + 1, 100)\n",
    "yy = -(w[0] * xx + b) / (w[1] + 1e-10)\n",
    "mask = (yy > X_2d[:, 1].min() - 2) & (yy < X_2d[:, 1].max() + 2)\n",
    "ax.plot(xx[mask], yy[mask], 'k--', alpha=0.4, linewidth=1, label='Probe boundary')\n",
    "\n",
    "ax.set_xlabel(f\"PC1 ({pca.explained_variance_ratio_[0]:.1%} variance)\")\n",
    "ax.set_ylabel(f\"PC2 ({pca.explained_variance_ratio_[1]:.1%} variance)\")\n",
    "ax.set_title(f\"Hidden States: Honest vs Sycophantic\\n(PCA of layer {mid_layer})\")\n",
    "ax.legend()\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"probe_separation.png\", dpi=150)\n",
    "plt.show()\n",
    "\n",
    "print(\"\\nBlue dots = honest framing. Red dots = sycophantic framing.\")\n",
    "print(\"The dashed line is where the linear probe separates them.\")\n",
    "print(\"The direction PERPENDICULAR to that line is the honesty vector.\")\n",
    "print(\"\\nSaved to probe_separation.png\")"
   ],
   "id": "cell-15"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7 — Extract the Steering Vector\n",
    "\n",
    "The probe's weight vector **is** the honesty direction. It's the line in\n",
    "the model's hidden space along which honest and sycophantic activations differ\n",
    "most. We normalize it to unit length so we can scale it later.\n",
    "\n",
    "**This is what you'll use in Notebook 2 to steer the model.**"
   ],
   "id": "cell-16"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# The probe's coefficient IS the concept direction\n",
    "honesty_direction = probe.coef_[0].copy()\n",
    "honesty_direction /= np.linalg.norm(honesty_direction)\n",
    "\n",
    "print(f\"Steering vector extracted.\")\n",
    "print(f\"  Shape: {honesty_direction.shape}\")\n",
    "print(f\"  Norm:  {np.linalg.norm(honesty_direction):.4f} (unit vector)\")\n",
    "print(f\"  Layer: {mid_layer}\")\n",
    "print(f\"\")\n",
    "print(f\"This vector IS what the model uses to represent honesty vs sycophancy.\")\n",
    "print(f\"Adding it to the model's activations will push it toward honesty.\")\n",
    "print(f\"Subtracting it will push toward sycophancy.\")\n",
    "\n",
    "# Save for Notebook 2\n",
    "np.save(\"honesty_direction.npy\", honesty_direction)\n",
    "print(f\"\\nSaved to honesty_direction.npy — you'll load this in Notebook 2.\")"
   ],
   "id": "cell-17"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8 — Quick Steering Preview\n",
    "\n",
    "Before we move to Notebook 2, let's do a quick sanity check: does adding this\n",
    "vector actually change the model's output?\n",
    "\n",
    "A unit vector is tiny next to the model's real activations, so we scale it by the\n",
    "typical size (norm) of a hidden state at this layer. `alpha` is then a fraction of that size."
   ],
   "id": "cell-18"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def make_hook(vector_tensor, alpha=1.0):\n",
    "    \"\"\"Forward hook that adds a steering vector to a layer's output.\"\"\"\n",
    "    def hook_fn(module, input, output):\n",
    "        if isinstance(output, tuple):\n",
    "            return (output[0] + alpha * vector_tensor,) + output[1:]\n",
    "        return output + alpha * vector_tensor\n",
    "    return hook_fn\n",
    "\n",
    "# Scale the unit direction to the typical hidden-state norm at this layer\n",
    "typical_norm = float(np.linalg.norm(X, axis=1).mean())\n",
    "print(f\"Typical hidden-state norm at layer {mid_layer}: {typical_norm:.1f}\\n\")\n",
    "vec_tensor = torch.tensor(honesty_direction * typical_norm, dtype=model.dtype).to(model.device)\n",
    "\n",
    "def _quick_generate(prompt, alpha=0.0):\n",
    "    messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "    # Qwen3.5 \"thinks\" before answering by default; turn that off for short answers\n",
    "    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,\n",
    "                                         enable_thinking=False)\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\").to(model.device)\n",
    "\n",
    "    handle = None\n",
    "    if alpha != 0:\n",
    "        handle = model.model.layers[mid_layer].register_forward_hook(make_hook(vec_tensor, alpha))\n",
    "    try:\n",
    "        with torch.no_grad():\n",
    "            out = model.generate(**inputs, max_new_tokens=100, temperature=0.7,\n",
    "                                 do_sample=True, pad_token_id=tokenizer.eos_token_id)\n",
    "    finally:\n",
    "        if handle:\n",
    "            handle.remove()\n",
    "\n",
    "    return tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True).strip()\n",
    "\n",
    "def quick_generate(prompt, alpha=0.0):\n",
    "    \"\"\"Generate with steering, logged to Langfuse as a generation.\"\"\"\n",
    "    with trace(\"steered-generation\", as_type=\"generation\", model=MODEL_NAME, input=prompt,\n",
    "               model_parameters={\"alpha\": alpha, \"layer\": mid_layer, \"temperature\": 0.7,\n",
    "                                 \"max_new_tokens\": 100}) as gen:\n",
    "        response = _quick_generate(prompt, alpha)\n",
    "        gen.update(output=response)\n",
    "    return response\n",
    "\n",
    "prompt = \"My friend's painting is terrible but they asked what I think. What should I say?\"\n",
    "print(f\"Prompt: {prompt}\\n\")\n",
    "\n",
    "for alpha, label in [(0, \"BASELINE\"), (0.2, f\"STEERED HONEST (+{0.2})\"), (-0.2, f\"STEERED SYCOPHANTIC (-{0.2})\")]:\n",
    "    print(f\"[{label}]\")\n",
    "    print(quick_generate(prompt, alpha))\n",
    "    print()"
   ],
   "id": "cell-19"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What You Just Did\n",
    "\n",
    "1. **Rented your own GPU** on vast.ai — not Colab, not a managed notebook. Your machine.\n",
    "2. **Loaded a 4-billion-parameter model** and looked at its internal structure\n",
    "3. **Created contrastive pairs** — the same questions framed as honest vs sycophantic\n",
    "4. **Extracted hidden states** — the model's internal representations at the middle layer\n",
    "5. **Trained a linear probe** — proved the model has a linearly separable concept of honesty\n",
    "6. **Visualized the separation** — saw the honesty direction in 2D\n",
    "7. **Extracted the steering vector** — the unit-length direction in activation space\n",
    "8. **Previewed steering** — saw the model's behavior change when you add/subtract the vector\n",
    "\n",
    "**The key insight:** honesty isn't stored in a single neuron. It's a *direction*\n",
    "across thousands of dimensions. And because it's a direction, you can move along it —\n",
    "making the model more honest, or more sycophantic, by adding or subtracting the vector.\n",
    "\n",
    "**Next: Notebook 2** — systematic steering experiments, adversarial red-teaming, and\n",
    "deploying this as an API on your own backend.\n",
    "\n",
    "**Don't forget:** stop or destroy your vast.ai instance when you're done — you're billed while it runs."
   ],
   "id": "cell-20"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Send any buffered traces to Langfuse\n",
    "if TRACING:\n",
    "    langfuse.flush()\n",
    "    print(f\"Traces flushed. Find this run in Langfuse under session {SESSION_ID}\")"
   ],
   "id": "cell-21"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}