{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook 2: A Probe for Every Concept\n",
    "\n",
    "**AI Alignment — The Multiverse School**\n",
    "\n",
    "In Notebook 1 you found *one* direction — honesty vs sycophancy — from 20 hand-written pairs.\n",
    "\n",
    "Now we scale up. You'll train a **separate linear probe for five concepts** a model might carry\n",
    "and act on without ever saying so out loud:\n",
    "\n",
    "| Concept | Data | Source |\n",
    "|---|---|---|\n",
    "| expansionism | 240 texts, matched pairs | synthetic |\n",
    "| colonialism | 240 texts, matched pairs | synthetic |\n",
    "| sexism | 240 synthetic pairs + 468 real posts | synthetic + ToxiGen |\n",
    "| racism | 2,724 real posts | ToxiGen (human-annotated) |\n",
    "| white supremacy | 2,374 real forum sentences | Stormfront corpus (human-annotated) |\n",
    "\n",
    "For each concept you'll answer three questions:\n",
    "1. **Is it there?** Can a linear probe detect it from the model's hidden states?\n",
    "2. **Where is it?** Which layer represents it most cleanly?\n",
    "3. **Is it its own thing?** Or is the \"racism\" probe secretly the same as the \"white supremacy\" probe?\n",
    "\n",
    "**Content warning:** the positive examples in this dataset are real hateful and prejudiced text,\n",
    "collected by researchers so that models like this one can be audited. We never print them in bulk.\n",
    "\n",
    "**What you need:** your vast.ai instance from Notebook 1, and the `data/probe_prompts/` folder\n",
    "uploaded next to this notebook."
   ],
   "id": "cell-00"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0 — Setup\n",
    "\n",
    "Same instance as Notebook 1. Upload the `data/` folder next to this notebook (drag it into Jupyter's file browser),\n",
    "so that `data/probe_prompts/all_concepts.jsonl` exists."
   ],
   "id": "cell-01"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import subprocess\n",
    "import sys\n",
    "\n",
    "subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"-U\",\n",
    "                       \"transformers>=5.17\", \"accelerate\", \"scikit-learn\", \"matplotlib\", \"pandas\"])\n",
    "\n",
    "import json\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "\n",
    "DATA = Path(\"data/probe_prompts/all_concepts.jsonl\")\n",
    "assert DATA.exists(), \"Upload the data/ folder next to this notebook first\"\n",
    "\n",
    "df = pd.DataFrame([json.loads(line) for line in DATA.read_text().splitlines()])\n",
    "# Where each row came from: \"synthetic\" pairs or a \"real\" annotated corpus\n",
    "df[\"origin\"] = np.where(df[\"format\"].isin([\"toxigen\", \"forum_post\"]), \"real\", \"synthetic\")\n",
    "\n",
    "print(f\"{len(df)} texts\\n\")\n",
    "print(df.groupby([\"concept\", \"origin\", \"label\"]).size().unstack().fillna(0).astype(int))"
   ],
   "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 = \"02_probe_every_concept\"\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": [
    "### What the labels mean\n",
    "\n",
    "- **label 1** — the text *endorses* the concept.\n",
    "- **label 0** — it doesn't. The negatives are chosen to be hard:\n",
    "  - *counter*: same topic, same format, same length, opposite stance (synthetic pairs)\n",
    "  - *discuss*: neutral or critical text **about** the concept — a history lecture on colonialism, a news report on a rally\n",
    "  - *benign_mention* / *same_forum_nonhate*: real posts about the same groups, or from the same forum, that aren't hateful\n",
    "\n",
    "Hard negatives matter. If every negative were a cooking recipe, the probe would learn \"is this about race?\" —\n",
    "not \"is this racist?\". Let's look at one matched pair:"
   ],
   "id": "cell-04"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "pair = df[df.pair_id == \"colonialism-p053\"].sort_values(\"label\", ascending=False)\n",
    "for _, r in pair.iterrows():\n",
    "    print(f\"[label {r.label} · {r.type}]\\n{r.text}\\n\")"
   ],
   "id": "cell-05"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 — Load the Model\n",
    "\n",
    "Same model as Notebook 1: **Qwen3.5-4B**."
   ],
   "id": "cell-06"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "MODEL_NAME = \"Qwen/Qwen3.5-4B\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16, device_map=\"auto\")\n",
    "model.eval()\n",
    "\n",
    "N_LAYERS = model.config.num_hidden_layers\n",
    "HIDDEN = model.config.hidden_size\n",
    "print(f\"{N_LAYERS} layers × {HIDDEN} dims · GPU memory used: {torch.cuda.memory_allocated() / 1e9:.1f} GB\")"
   ],
   "id": "cell-07"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 — Extract Hidden States from Every Layer\n",
    "\n",
    "In Notebook 1 we took the **last token** at the **middle layer**. This time:\n",
    "\n",
    "- **Every layer.** We don't know in advance where each concept lives, so we record all of them and let the data tell us.\n",
    "- **Mean over tokens.** A racist sentence can put its payload anywhere, not just at the end. Averaging the hidden\n",
    "  state over all tokens gives a summary of the whole text.\n",
    "\n",
    "That's 6,000+ texts × 33 hidden states × 2,560 dims. We batch the texts and store float16 to keep memory reasonable (~1 GB)."
   ],
   "id": "cell-08"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "tokenizer.padding_side = \"right\"  # padding after the text never affects earlier tokens (causal model)\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "@torch.no_grad()\n",
    "def extract_all_layers(texts, batch_size=32, max_length=256):\n",
    "    \"\"\"Return float16 array [n_texts, n_layers + 1, hidden]: mean-pooled hidden state per layer.\"\"\"\n",
    "    out = np.zeros((len(texts), N_LAYERS + 1, HIDDEN), dtype=np.float16)\n",
    "    for start in range(0, len(texts), batch_size):\n",
    "        batch = texts[start:start + batch_size]\n",
    "        enc = tokenizer(batch, return_tensors=\"pt\", padding=True, truncation=True,\n",
    "                        max_length=max_length).to(model.device)\n",
    "        hs = model(**enc, output_hidden_states=True).hidden_states          # tuple of [B, T, H]\n",
    "        mask = enc[\"attention_mask\"].unsqueeze(-1).to(hs[0].dtype)          # [B, T, 1]\n",
    "        pooled = torch.stack([(h * mask).sum(1) / mask.sum(1) for h in hs], dim=1)  # [B, L+1, H]\n",
    "        out[start:start + len(batch)] = pooled.float().cpu().numpy()\n",
    "        if len(texts) > batch_size and (start // batch_size) % 40 == 0:\n",
    "            print(f\"  {start + len(batch)}/{len(texts)}\")\n",
    "    return out\n",
    "\n",
    "FEATURES = Path(\"hidden_states_all_concepts.npy\")\n",
    "if FEATURES.exists():\n",
    "    H = np.load(FEATURES)\n",
    "    print(\"Loaded cached hidden states.\")\n",
    "else:\n",
    "    with trace(\"extract-hidden-states\", input={\"n_texts\": len(df), \"layers\": N_LAYERS + 1},\n",
    "               metadata={\"model\": MODEL_NAME, \"pooling\": \"mean\"}) as obs:\n",
    "        H = extract_all_layers(df.text.tolist())\n",
    "        obs.update(output={\"shape\": list(H.shape)})\n",
    "    np.save(FEATURES, H)\n",
    "print(f\"\\nHidden states: {H.shape}  (texts, layers, dims)\")"
   ],
   "id": "cell-09"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 — Train a Probe per Concept, at Every Layer\n",
    "\n",
    "For each concept, at each layer, we train a logistic-regression probe and measure it on held-out texts.\n",
    "\n",
    "Two details that keep us honest:\n",
    "- **Pairs stay together.** Both halves of a matched pair go to the same split. Otherwise the probe could\n",
    "  memorize \"I saw the other half of this pair in training\" instead of learning the concept.\n",
    "- **We score with AUC** (area under the ROC curve): 0.5 = coin flip, 1.0 = perfect ranking. It doesn't depend on\n",
    "  picking a threshold."
   ],
   "id": "cell-10"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import accuracy_score, roc_auc_score\n",
    "from sklearn.model_selection import GroupShuffleSplit\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "CONCEPTS = [\"expansionism\", \"colonialism\", \"sexism\", \"racism\", \"white_supremacy\"]\n",
    "\n",
    "def make_probe():\n",
    "    return make_pipeline(StandardScaler(), LogisticRegression(C=0.05, max_iter=2000))\n",
    "\n",
    "# Split once per concept: 60% train, 20% validation (to pick the layer), 20% test (reported)\n",
    "splits = {}\n",
    "for concept in CONCEPTS:\n",
    "    idx = np.where(df.concept == concept)[0]\n",
    "    groups = df.pair_id.iloc[idx].fillna(df.id.iloc[idx]).values   # pairs share a group\n",
    "    trainval, test = next(GroupShuffleSplit(1, test_size=0.2, random_state=0).split(idx, groups=groups))\n",
    "    train, val = next(GroupShuffleSplit(1, test_size=0.25, random_state=0).split(trainval, groups=groups[trainval]))\n",
    "    splits[concept] = {\"train\": idx[trainval[train]], \"val\": idx[trainval[val]], \"test\": idx[test]}\n",
    "    print(f\"{concept:16s} train {len(splits[concept]['train']):5d} · val {len(splits[concept]['val']):4d} · test {len(splits[concept]['test']):4d}\")\n",
    "\n",
    "y_all = df.label.values\n",
    "layer_auc = {c: [] for c in CONCEPTS}\n",
    "for concept in CONCEPTS:\n",
    "    s = splits[concept]\n",
    "    for layer in range(N_LAYERS + 1):\n",
    "        probe = make_probe().fit(H[s[\"train\"], layer].astype(np.float32), y_all[s[\"train\"]])\n",
    "        scores = probe.predict_proba(H[s[\"val\"], layer].astype(np.float32))[:, 1]\n",
    "        layer_auc[concept].append(roc_auc_score(y_all[s[\"val\"]], scores))\n",
    "    best = int(np.argmax(layer_auc[concept]))\n",
    "    print(f\"{concept:16s} best layer {best:2d}  (val AUC {layer_auc[concept][best]:.3f})\")\n",
    "\n",
    "with trace(\"layer-sweep\", input={\"concepts\": CONCEPTS}, metadata={\"model\": MODEL_NAME}) as obs:\n",
    "    obs.update(output={c: {\"best_layer\": int(np.argmax(v)), \"val_auc_by_layer\": [round(a, 4) for a in v]}\n",
    "                       for c, v in layer_auc.items()})"
   ],
   "id": "cell-11"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4 — Where Does Each Concept Live?\n",
    "\n",
    "Each line is one concept. Layer 0 is the raw token embeddings — before the model has \"thought\" at all.\n",
    "Watch how quickly each line climbs, and where it peaks."
   ],
   "id": "cell-12"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "for concept in CONCEPTS:\n",
    "    ax.plot(range(N_LAYERS + 1), layer_auc[concept], marker=\"o\", markersize=3, label=concept)\n",
    "ax.axhline(0.5, color=\"gray\", linestyle=\":\", linewidth=1)\n",
    "ax.set_xlabel(\"Layer (0 = embeddings)\")\n",
    "ax.set_ylabel(\"Validation AUC\")\n",
    "ax.set_title(\"How well a linear probe detects each concept, layer by layer\")\n",
    "ax.set_ylim(0.45, 1.01)\n",
    "ax.legend()\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"probe_layer_sweep.png\", dpi=150)\n",
    "plt.show()"
   ],
   "id": "cell-13"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Things to notice:**\n",
    "- Layer 0 is often already well above 0.5. Some signal is in the *words themselves* (a slur, \"annex\", \"natives\").\n",
    "  The interesting part is how much the model *adds* on top of that in the middle layers.\n",
    "- The synthetic concepts (expansionism, colonialism) were built as matched pairs where the words barely differ —\n",
    "  so they need the model's understanding of *stance*, not vocabulary."
   ],
   "id": "cell-14"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 — Final Probes: Train at the Best Layer, Report on Test\n",
    "\n",
    "Now we retrain each probe on train + validation at its best layer and measure it **once** on the test set it has never seen."
   ],
   "id": "cell-15"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "final = {}\n",
    "rows = []\n",
    "for concept in CONCEPTS:\n",
    "    s = splits[concept]\n",
    "    layer = int(np.argmax(layer_auc[concept]))\n",
    "    trainval = np.concatenate([s[\"train\"], s[\"val\"]])\n",
    "    probe = make_probe().fit(H[trainval, layer].astype(np.float32), y_all[trainval])\n",
    "    scores = probe.predict_proba(H[s[\"test\"], layer].astype(np.float32))[:, 1]\n",
    "    final[concept] = {\"probe\": probe, \"layer\": layer}\n",
    "    rows.append({\"concept\": concept, \"layer\": layer,\n",
    "                 \"test AUC\": roc_auc_score(y_all[s[\"test\"]], scores),\n",
    "                 \"test accuracy\": accuracy_score(y_all[s[\"test\"]], scores > 0.5),\n",
    "                 \"n test\": len(s[\"test\"])})\n",
    "\n",
    "results = pd.DataFrame(rows).set_index(\"concept\")\n",
    "with trace(\"final-probes\", input={\"concepts\": CONCEPTS}, metadata={\"model\": MODEL_NAME}) as obs:\n",
    "    obs.update(output=results.reset_index().to_dict(\"records\"))\n",
    "    for concept, r in results.iterrows():\n",
    "        obs.score(name=f\"test_auc_{concept}\", value=float(r[\"test AUC\"]))\n",
    "results.style.format({\"test AUC\": \"{:.3f}\", \"test accuracy\": \"{:.1%}\"})"
   ],
   "id": "cell-16"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Which negatives fool the probe?\n",
    "\n",
    "Accuracy broken down by the *type* of text. If the probe does well on `counter` but badly on `discuss`,\n",
    "it's learned \"talks about colonialism\" rather than \"endorses colonialism\"."
   ],
   "id": "cell-17"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "breakdown = []\n",
    "for concept in CONCEPTS:\n",
    "    s, f = splits[concept], final[concept]\n",
    "    test = df.iloc[s[\"test\"]].copy()\n",
    "    test[\"pred\"] = f[\"probe\"].predict(H[s[\"test\"], f[\"layer\"]].astype(np.float32))\n",
    "    for t, g in test.groupby(\"type\"):\n",
    "        breakdown.append({\"concept\": concept, \"type\": t, \"n\": len(g), \"accuracy\": (g.pred == g.label).mean()})\n",
    "pd.DataFrame(breakdown).pivot(index=\"concept\", columns=\"type\", values=\"accuracy\").style.format(\"{:.0%}\", na_rep=\"\")"
   ],
   "id": "cell-18"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6 — Is Each Concept Its Own Thing?\n",
    "\n",
    "Two ways to check whether our five probes found five *different* directions:\n",
    "\n",
    "1. **Cross-detection.** Run every probe over every concept's positive examples. Does the racism probe fire on\n",
    "   white-supremacist sentences? Does the colonialism probe fire on expansionism?\n",
    "2. **Cosine similarity** between probe directions that live at the same layer. 1.0 = same direction, 0 = unrelated."
   ],
   "id": "cell-19"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import matplotlib.colors as mcolors\n",
    "\n",
    "# Every probe scores every concept's held-out *positive* texts\n",
    "cross = np.zeros((len(CONCEPTS), len(CONCEPTS)))\n",
    "for i, probe_c in enumerate(CONCEPTS):\n",
    "    f = final[probe_c]\n",
    "    for j, data_c in enumerate(CONCEPTS):\n",
    "        pos = splits[data_c][\"test\"][y_all[splits[data_c][\"test\"]] == 1]\n",
    "        cross[i, j] = f[\"probe\"].predict_proba(H[pos, f[\"layer\"]].astype(np.float32))[:, 1].mean()\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 6))\n",
    "im = ax.imshow(cross, cmap=\"Blues\", vmin=0, vmax=1)\n",
    "ax.set_xticks(range(len(CONCEPTS)), CONCEPTS, rotation=30, ha=\"right\")\n",
    "ax.set_yticks(range(len(CONCEPTS)), [f\"{c} probe\" for c in CONCEPTS])\n",
    "ax.set_xlabel(\"Positive examples of…\")\n",
    "for i in range(len(CONCEPTS)):\n",
    "    for j in range(len(CONCEPTS)):\n",
    "        ax.text(j, i, f\"{cross[i, j]:.2f}\", ha=\"center\", va=\"center\",\n",
    "                color=\"white\" if cross[i, j] > 0.6 else \"black\")\n",
    "ax.set_title(\"Mean probe score on each concept's positives\")\n",
    "plt.colorbar(im, fraction=0.046)\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"probe_cross_detection.png\", dpi=150)\n",
    "plt.show()"
   ],
   "id": "cell-20"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Cosine similarity between raw-space probe directions, all compared at one shared layer\n",
    "SHARED_LAYER = int(np.median([final[c][\"layer\"] for c in CONCEPTS]))\n",
    "\n",
    "def direction_at(concept, layer):\n",
    "    s = splits[concept]\n",
    "    trainval = np.concatenate([s[\"train\"], s[\"val\"]])\n",
    "    p = make_probe().fit(H[trainval, layer].astype(np.float32), y_all[trainval])\n",
    "    scaler, lr = p.named_steps[\"standardscaler\"], p.named_steps[\"logisticregression\"]\n",
    "    d = lr.coef_[0] / scaler.scale_      # undo the standardization: a direction in the model's own space\n",
    "    return d / np.linalg.norm(d)\n",
    "\n",
    "dirs = np.stack([direction_at(c, SHARED_LAYER) for c in CONCEPTS])\n",
    "cos = dirs @ dirs.T\n",
    "print(f\"Cosine similarity between probe directions at layer {SHARED_LAYER}:\\n\")\n",
    "print(pd.DataFrame(cos, index=CONCEPTS, columns=CONCEPTS).round(2))"
   ],
   "id": "cell-21"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Discussion questions:**\n",
    "- Which pairs of concepts overlap the most? Is that a flaw in the probes, or a true fact about how these ideas are related?\n",
    "- White supremacy and racism come from different datasets (a white-nationalist forum vs. ToxiGen posts). If their probes\n",
    "  overlap, is that because the *ideas* overlap — or because both datasets share the word \"white\"?\n",
    "- Colonialism and expansionism were written to be distinct. Did the model keep them apart?"
   ],
   "id": "cell-22"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7 — Does the Sexism Probe Generalize?\n",
    "\n",
    "Our sexism data has two very different sources: 240 synthetic matched pairs, and 468 real ToxiGen posts.\n",
    "Train on one, test on the other. If the probe learned *sexism* — and not the writing style of one dataset —\n",
    "it should transfer."
   ],
   "id": "cell-23"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "layer = final[\"sexism\"][\"layer\"]\n",
    "sx = df[df.concept == \"sexism\"]\n",
    "transfer = []\n",
    "for train_origin, test_origin in [(\"synthetic\", \"real\"), (\"real\", \"synthetic\")]:\n",
    "    tr = sx.index[sx.origin == train_origin]\n",
    "    te = sx.index[sx.origin == test_origin]\n",
    "    p = make_probe().fit(H[tr, layer].astype(np.float32), y_all[tr])\n",
    "    auc = roc_auc_score(y_all[te], p.predict_proba(H[te, layer].astype(np.float32))[:, 1])\n",
    "    transfer.append({\"train on\": train_origin, \"test on\": test_origin, \"AUC\": auc})\n",
    "with trace(\"sexism-transfer\", metadata={\"layer\": layer}) as obs:\n",
    "    obs.update(output=transfer)\n",
    "pd.DataFrame(transfer).style.format({\"AUC\": \"{:.3f}\"})"
   ],
   "id": "cell-24"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8 — Try It Yourself\n",
    "\n",
    "Score any text with all five probes. Try subtle cases: a historical quote, a joke, a sentence that\n",
    "*criticizes* colonialism, a sentence that's prejudiced but polite."
   ],
   "id": "cell-25"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def score_text(text):\n",
    "    with trace(\"score-text\", input=text, metadata={\"model\": MODEL_NAME}) as obs:\n",
    "        h = extract_all_layers([text])[0]\n",
    "        scores = {c: float(final[c][\"probe\"].predict_proba(h[final[c][\"layer\"]].astype(np.float32)[None])[0, 1])\n",
    "                  for c in CONCEPTS}\n",
    "        obs.update(output=scores)\n",
    "    return scores\n",
    "\n",
    "examples = [\n",
    "    \"The river valley was always ours; it's time our borders reflected that.\",\n",
    "    \"Historians still debate how much the railways built under British rule benefited ordinary Indians.\",\n",
    "    \"Honestly she's pretty good at engineering, for a woman.\",\n",
    "    \"Our neighborhood association welcomes families from every background.\",\n",
    "]\n",
    "for text in examples:\n",
    "    print(text)\n",
    "    print(\"   \" + \"  \".join(f\"{c}: {p:.2f}\" for c, p in score_text(text).items()) + \"\\n\")"
   ],
   "id": "cell-26"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Look closely at the results.** Probes that score 0.95+ AUC on the test set can still miss, or misfire on, a single\n",
    "sentence you write yourself. \"Good at engineering, for a woman\" is condescending, polite sexism, which is rare in the\n",
    "training data. Which probe fires on it, and why might that be? Probes only know what their data showed them.\n",
    "Write five more hard cases and see where they break."
   ],
   "id": "cell-27"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9 — Save the Probes\n",
    "\n",
    "Each probe is saved as a **direction in the model's own activation space** plus the layer it lives at —\n",
    "the same form as `honesty_direction.npy` from Notebook 1. That's what you need to *monitor* a model for\n",
    "these concepts, or to *steer* away from them."
   ],
   "id": "cell-28"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "out_dir = Path(\"probes\")\n",
    "out_dir.mkdir(exist_ok=True)\n",
    "for concept in CONCEPTS:\n",
    "    f = final[concept]\n",
    "    scaler, lr = f[\"probe\"].named_steps[\"standardscaler\"], f[\"probe\"].named_steps[\"logisticregression\"]\n",
    "    raw_w = lr.coef_[0] / scaler.scale_                       # weights on raw (unstandardized) hidden states\n",
    "    raw_b = lr.intercept_[0] - (scaler.mean_ / scaler.scale_) @ lr.coef_[0]\n",
    "    np.savez(out_dir / f\"{concept}.npz\",\n",
    "             direction=(raw_w / np.linalg.norm(raw_w)).astype(np.float32),   # unit vector, for steering\n",
    "             weights=raw_w.astype(np.float32), bias=np.float32(raw_b),        # logit = h @ weights + bias\n",
    "             layer=f[\"layer\"], pooling=\"mean\", model=MODEL_NAME,\n",
    "             test_auc=results.loc[concept, \"test AUC\"])\n",
    "    print(f\"saved probes/{concept}.npz  (layer {f['layer']}, test AUC {results.loc[concept, 'test AUC']:.3f})\")\n",
    "\n",
    "# Sanity check: the raw-space weights reproduce the sklearn probe exactly\n",
    "c = CONCEPTS[0]\n",
    "saved = np.load(out_dir / f\"{c}.npz\")\n",
    "h = H[splits[c][\"test\"][:5], int(saved[\"layer\"])].astype(np.float32)\n",
    "manual = 1 / (1 + np.exp(-(h @ saved[\"weights\"] + saved[\"bias\"])))\n",
    "assert np.allclose(manual, final[c][\"probe\"].predict_proba(h)[:, 1], atol=1e-4)\n",
    "print(\"\\nRaw-space weights verified against the sklearn probe.\")"
   ],
   "id": "cell-29"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What You Just Did\n",
    "\n",
    "1. **Scaled up from one concept to five**, using a mix of synthetic matched pairs and real, human-annotated data\n",
    "2. **Swept every layer** to find where each concept is represented most cleanly\n",
    "3. **Trained and tested five linear probes** on held-out data, keeping matched pairs together\n",
    "4. **Checked the hard negatives** — can the probe tell *endorsing* an idea from *talking about* it?\n",
    "5. **Measured overlap** between concepts — both behaviorally (cross-detection) and geometrically (cosine similarity)\n",
    "6. **Tested generalization** across data sources\n",
    "7. **Saved each probe** as a direction in activation space\n",
    "\n",
    "**The key insight:** a model doesn't have to *say* something racist for you to detect that it's representing racism.\n",
    "These probes read the model's internal state directly — which is exactly what you want when auditing a model\n",
    "that might have learned to hide what it's thinking.\n",
    "\n",
    "**Don't forget:** stop or destroy your vast.ai instance when you're done.\n",
    "\n",
    "---\n",
    "*Data sources:* ToxiGen (Hartvigsen et al., 2022). Stormfront hate speech corpus (de Gibert et al., 2018),\n",
    "CC BY-SA 3.0 ES — derived data must carry the same license."
   ],
   "id": "cell-30"
  },
  {
   "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-31"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}