From 5bc6f27c77f64f05142bc37c740d452cd9e698b3 Mon Sep 17 00:00:00 2001
From: Theia Vogel <theia@vgel.me>
Date: Tue, 23 Sep 2025 18:21:19 -0700
Subject: [PATCH 1/2] add hook for custom compute_hiddens

---
 CHANGELOG                   |   2 +
 notebooks/model_delta.ipynb | 264 ++++++++++++++++++++++++++++++++++++
 repeng/extract.py           |  31 ++++-
 repeng/tests.py             |  22 +++
 4 files changed, 316 insertions(+), 3 deletions(-)
 create mode 100644 notebooks/model_delta.ipynb

diff --git a/CHANGELOG b/CHANGELOG
index 2baeaff..f758493 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -7,6 +7,8 @@
 * Fix for Qwen2 inference `attention_type`. (#73)
 * Reasoning dataset for steering thinking models. (#69)
   * Thanks to @wassname ! :tada:
+* Add `compute_hiddens` hook for fancy things. (#74)
+  * See `notebooks/model_delta.ipynb` for an example.
 
 ## 0.4.0 - 2024-12-13
 
diff --git a/notebooks/model_delta.ipynb b/notebooks/model_delta.ipynb
new file mode 100644
index 0000000..1bb124f
--- /dev/null
+++ b/notebooks/model_delta.ipynb
@@ -0,0 +1,264 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "6ae9b6f2-30f7-4a33-b6f4-fc9ebe358598",
+   "metadata": {},
+   "source": [
+    "Here's an example of training a vector on the difference of the same prompt between two models, instead of the difference of two prompts on the same model.\n",
+    "\n",
+    "Needs `datasets`: `pip install datasets`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "78e6c8e2-8ac8-4f9d-85cf-3c75ec20642e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import datasets\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "from repeng import ControlModel, ControlVector, DatasetEntry\n",
+    "from repeng.extract import batched_get_hiddens"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "f1c5d67f-06b9-4abd-8403-80431c4fb2c0",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5189cf5e6cd54cf180998c206307eb4c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0b2ff2b38ba34528a3d2abbe1f8c31f0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "MODEL_A = \"Qwen/Qwen2.5-7B\"\n",
+    "MODEL_B = \"Qwen/Qwen2.5-7B-Instruct\"\n",
+    "DATASET = \"agentlans/wikipedia-paragraphs\"\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(MODEL_A)\n",
+    "model_a = AutoModelForCausalLM.from_pretrained(\n",
+    "    MODEL_A, dtype=torch.bfloat16, device_map=\"cuda\"\n",
+    ")\n",
+    "model_b = AutoModelForCausalLM.from_pretrained(\n",
+    "    MODEL_B, dtype=torch.bfloat16, device_map=\"cuda\"\n",
+    ")\n",
+    "dataset = datasets.load_dataset(DATASET)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "6ee39f5b-96f5-4f77-a865-7008467c0e81",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_dataset = []\n",
+    "for text in dataset[\"train\"].take(100)[\"text\"]:\n",
+    "    text = \" \".join(text.split(\" \")[:100])\n",
+    "    chat = tokenizer.apply_chat_template(\n",
+    "        [{\"role\": \"user\", \"content\": text}], tokenize=False, add_generation_prompt=True\n",
+    "    )\n",
+    "    train_dataset.append(DatasetEntry(positive=text, negative=chat))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "4431f11b-c7d0-43ee-84e7-86d08bb6dac1",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Hooked compute_hiddens\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████| 4/4 [00:01<00:00,  2.95it/s]\n",
+      "100%|█████████████████████████████████████████████| 4/4 [00:01<00:00,  2.81it/s]\n",
+      "100%|██████████████████████████████████████████| 27/27 [00:00<00:00, 124.30it/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# we get passed model, tokenizer, ... from ControlVector.train\n",
+    "# we don't need these, so ignore them with **kwargs\n",
+    "def compute_hiddens(train_strs, hidden_layers, batch_size, **kwargs):\n",
+    "    print(\"Hooked compute_hiddens\")\n",
+    "\n",
+    "    a_train_strs, b_train_strs = train_strs[::2], train_strs[1::2]\n",
+    "    assert len(a_train_strs) == len(b_train_strs)\n",
+    "\n",
+    "    a_hiddens = batched_get_hiddens(\n",
+    "        model_a, tokenizer, a_train_strs, hidden_layers, batch_size\n",
+    "    )\n",
+    "    b_hiddens = batched_get_hiddens(\n",
+    "        model_b, tokenizer, b_train_strs, hidden_layers, batch_size\n",
+    "    )\n",
+    "    interleaved = {}\n",
+    "    for layer in hidden_layers:\n",
+    "        ah, bh = a_hiddens[layer], b_hiddens[layer]\n",
+    "        i = np.stack((ah, bh))\n",
+    "        i = i.transpose(1, 0, *range(2, i.ndim))\n",
+    "        i = i.reshape((ah.shape[0] + bh.shape[0], ah.shape[1]))  # ex*2, hidden_dim\n",
+    "        interleaved[layer] = i\n",
+    "    return interleaved\n",
+    "\n",
+    "\n",
+    "completion_vector = ControlVector.train(\n",
+    "    model=model_a,\n",
+    "    tokenizer=tokenizer,\n",
+    "    dataset=train_dataset,\n",
+    "    compute_hiddens=compute_hiddens,\n",
+    "    method=\"pca_center\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "230a4ae7-8488-44e2-a860-7eb71d518471",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "# baseline:\n",
+      "Hurt-Proofing Your Child\n",
+      "By: Dr. Michael R"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/ubuntu/repeng/repeng/control.py:37: UserWarning: Trying to rewrap a wrapped model! Probably not what you want! Try calling .unwrap first.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      ". Thompson\n",
+      "The most important thing you can do to help your child is to help him or her develop a sense of self-worth. This is the most important ingredient in a child's ability to cope with the world. A child who feels good about himself or herself is more likely to be able to handle the challenges of life. A child who feels good about himself or herself is more likely to be able to handle the challenges of life. A child who feels good about himself or herself is more likely to be able to handle the challenges of life. A child who feels good about\n",
+      "\n",
+      "# steered towards instruct\n",
+      "Hurt-Proofing Your Child: How to Help Your Child Build Resilience\n",
+      "Resilience is the ability to bounce back from adversity. It’s the ability to adapt to change and to learn from mistakes.It’s the ability to cope with stress and to find meaning in difficult situations.\n",
+      "Resilience is not something that you’re born with. It’s something that you can develop over time.\n",
+      "Here are some tips to help your child build resilience:\n",
+      "1. Teach your child to identify and label their emotions. This will help them understand what they’re feeling and why they’re feeling it.\n",
+      "2. Encourage your child to talk about their\n",
+      "\n",
+      "# steered away from instruct\n",
+      "Hurt-owning and --owning households in the United States are more likely to be in poverty than those who own no property. This is true for all races and ethnicities, and for all income levels. The gap between the property-owning and non-owning poor is greatest for African Americans and Hispanics. The gap is also greatest for the poorest households. The gap between the property-owning and non- owning poor is greatest for African Americans and Hispanics. The gap is also greatest for the poorest households. The gap between the property-owning and non- owning poor is greatest for African Americans and Hispanics. The gap\n"
+     ]
+    }
+   ],
+   "source": [
+    "from transformers import TextStreamer\n",
+    "\n",
+    "\n",
+    "class TokenStreamer(TextStreamer):\n",
+    "    def _is_chinese_char(*args, **kwargs):\n",
+    "        return True\n",
+    "\n",
+    "\n",
+    "def generate_with_vector(\n",
+    "    prompt: str,\n",
+    "    vectors,\n",
+    "    model=model_a,\n",
+    "    max_new_tokens: int = 128,\n",
+    "):\n",
+    "    ctl = ControlModel(model, list(range(1, 28)))\n",
+    "    input_ids = tokenizer(prompt, return_tensors=\"pt\")\n",
+    "    settings = {\n",
+    "        \"pad_token_id\": tokenizer.eos_token_id,  # silence warning\n",
+    "        \"do_sample\": False,  # temperature=0\n",
+    "        \"max_new_tokens\": max_new_tokens,\n",
+    "    }\n",
+    "\n",
+    "    def generate():\n",
+    "        ctl.generate(\n",
+    "            streamer=TokenStreamer(tokenizer), **input_ids.to(ctl.device), **settings\n",
+    "        )\n",
+    "\n",
+    "    ctl.reset()\n",
+    "    print(\"# baseline:\")\n",
+    "    generate()\n",
+    "    for label, v in vectors:\n",
+    "        print(f\"\\n# {label}\")\n",
+    "        ctl.set_control(v)\n",
+    "        generate()\n",
+    "    ctl.reset()\n",
+    "    ctl.unwrap()\n",
+    "\n",
+    "\n",
+    "generate_with_vector(\n",
+    "    \"Hurt-\",\n",
+    "    [\n",
+    "        (\"steered towards instruct\", completion_vector * -1.5),\n",
+    "        (\"steered away from instruct\", completion_vector * 2.0),\n",
+    "    ],\n",
+    ")"
+   ]
+  }
+ ],
+ "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.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/repeng/extract.py b/repeng/extract.py
index 8d78a75..80cb541 100644
--- a/repeng/extract.py
+++ b/repeng/extract.py
@@ -45,6 +45,10 @@ def train(
                     Defaults to 32. Try reducing this if you're running out of memory.
                 method (str, optional): The training method to use. Can be either
                     "pca_diff" or "pca_center". Defaults to "pca_diff".
+                compute_hiddens (Callable, optional): Override hidden state computation.
+                    See signature of `read_representations`.
+                transform_hiddens (Callable, optional): Transform the hidden states after
+                    they are computed. See signature of `read_representations`.
 
         Returns:
             ControlVector: The trained vector.
@@ -240,6 +244,17 @@ def __truediv__(self, other: int | float | np.number) -> "ControlVector":
         return self.__mul__(1 / other)
 
 
+class ComputeHiddens(typing.Protocol):
+    def __call__(
+        self,
+        model: "PreTrainedModel | ControlModel",
+        tokenizer: PreTrainedTokenizerBase,
+        train_strs: list[str],
+        hidden_layers: list[int],
+        batch_size: int,
+    ) -> dict[int, np.ndarray]: ...
+
+
 def read_representations(
     model: "PreTrainedModel | ControlModel",
     tokenizer: PreTrainedTokenizerBase,
@@ -247,6 +262,7 @@ def read_representations(
     hidden_layers: typing.Iterable[int] | None = None,
     batch_size: int = 32,
     method: typing.Literal["pca_diff", "pca_center", "umap"] = "pca_diff",
+    compute_hiddens: ComputeHiddens | None = None,
     transform_hiddens: (
         typing.Callable[[dict[int, np.ndarray]], dict[int, np.ndarray]] | None
     ) = None,
@@ -264,9 +280,18 @@ def read_representations(
     # the order is [positive, negative, positive, negative, ...]
     train_strs = [s for ex in inputs for s in (ex.positive, ex.negative)]
 
-    layer_hiddens = batched_get_hiddens(
-        model, tokenizer, train_strs, hidden_layers, batch_size
-    )
+    if compute_hiddens is None:
+        layer_hiddens = batched_get_hiddens(
+            model, tokenizer, train_strs, hidden_layers, batch_size
+        )
+    else:
+        layer_hiddens = compute_hiddens(
+            model=model,
+            tokenizer=tokenizer,
+            train_strs=train_strs,
+            hidden_layers=hidden_layers,
+            batch_size=batch_size,
+        )
 
     if transform_hiddens is not None:
         layer_hiddens = transform_hiddens(layer_hiddens)
diff --git a/repeng/tests.py b/repeng/tests.py
index 94ea315..d5b4782 100644
--- a/repeng/tests.py
+++ b/repeng/tests.py
@@ -3,11 +3,13 @@
 import pathlib
 import tempfile
 
+import numpy as np
 import pytest
 from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerBase
 
 from . import ControlModel, ControlVector, DatasetEntry
 from .control import model_layer_list
+from .extract import batched_get_hiddens
 
 
 @pytest.mark.slow
@@ -125,6 +127,26 @@ def test_layer_list_real():
     assert len(model_layer_list(lts)) == 4
 
 
+@pytest.mark.slow
+@pytest.mark.filterwarnings("ignore:invalid value")  # all-zeros is degenerate
+def test_hook_compute_hiddens():
+    tokenizer, model = load_llama_tinystories_model()
+    suffixes = load_suffixes()[:2]
+    dataset = make_dataset("{persona}", ["a"], ["b"], suffixes)
+
+    def compute_hiddens(model, tokenizer, train_strs, hidden_layers, batch_size):
+        h = batched_get_hiddens(model, tokenizer, train_strs, hidden_layers, batch_size)
+        return {k: np.zeros_like(v) for k, v in h.items()}
+
+    cvec = ControlVector.train(
+        model, tokenizer, dataset, compute_hiddens=compute_hiddens
+    )
+    assert len(cvec.directions) == 3
+    for v in cvec.directions.values():
+        assert v[0] == 1.0
+        assert (v[1:] == 0.0).all()
+
+
 def test_layer_list_override():
     import torch
     from transformers.models.llama import LlamaForCausalLM, LlamaConfig

From 632e3b0033e204f33559849aff6b0838aecba6ea Mon Sep 17 00:00:00 2001
From: Theia Vogel <theia@vgel.me>
Date: Tue, 23 Sep 2025 18:27:30 -0700
Subject: [PATCH 2/2] fix notebook

---
 notebooks/model_delta.ipynb | 42 ++++++++++++-------------------------
 1 file changed, 13 insertions(+), 29 deletions(-)

diff --git a/notebooks/model_delta.ipynb b/notebooks/model_delta.ipynb
index 1bb124f..0775a2f 100644
--- a/notebooks/model_delta.ipynb
+++ b/notebooks/model_delta.ipynb
@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 2,
    "id": "78e6c8e2-8ac8-4f9d-85cf-3c75ec20642e",
    "metadata": {},
    "outputs": [],
@@ -34,7 +34,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5189cf5e6cd54cf180998c206307eb4c",
+       "model_id": "8bbb21e1a2ae456aacdf88b9af4c94d6",
        "version_major": 2,
        "version_minor": 0
       },
@@ -48,7 +48,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0b2ff2b38ba34528a3d2abbe1f8c31f0",
+       "model_id": "08e53b2b34a247aea7a5a42b9acf3b28",
        "version_major": 2,
        "version_minor": 0
       },
@@ -93,7 +93,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 5,
    "id": "4431f11b-c7d0-43ee-84e7-86d08bb6dac1",
    "metadata": {},
    "outputs": [
@@ -108,9 +108,9 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|█████████████████████████████████████████████| 4/4 [00:01<00:00,  2.95it/s]\n",
-      "100%|█████████████████████████████████████████████| 4/4 [00:01<00:00,  2.81it/s]\n",
-      "100%|██████████████████████████████████████████| 27/27 [00:00<00:00, 124.30it/s]\n"
+      "100%|█████████████████████████████████████████████| 4/4 [00:01<00:00,  2.35it/s]\n",
+      "100%|█████████████████████████████████████████████| 4/4 [00:01<00:00,  2.48it/s]\n",
+      "100%|███████████████████████████████████████████| 27/27 [00:00<00:00, 57.41it/s]\n"
      ]
     }
    ],
@@ -150,7 +150,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 37,
+   "execution_count": 6,
    "id": "230a4ae7-8488-44e2-a860-7eb71d518471",
    "metadata": {},
    "outputs": [
@@ -160,31 +160,15 @@
      "text": [
       "# baseline:\n",
       "Hurt-Proofing Your Child\n",
-      "By: Dr. Michael R"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/ubuntu/repeng/repeng/control.py:37: UserWarning: Trying to rewrap a wrapped model! Probably not what you want! Try calling .unwrap first.\n",
-      "  warnings.warn(\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      ". Thompson\n",
+      "By: Dr. Michael R. Thompson\n",
       "The most important thing you can do to help your child is to help him or her develop a sense of self-worth. This is the most important ingredient in a child's ability to cope with the world. A child who feels good about himself or herself is more likely to be able to handle the challenges of life. A child who feels good about himself or herself is more likely to be able to handle the challenges of life. A child who feels good about himself or herself is more likely to be able to handle the challenges of life. A child who feels good about\n",
       "\n",
       "# steered towards instruct\n",
       "Hurt-Proofing Your Child: How to Help Your Child Build Resilience\n",
-      "Resilience is the ability to bounce back from adversity. It’s the ability to adapt to change and to learn from mistakes.It’s the ability to cope with stress and to find meaning in difficult situations.\n",
-      "Resilience is not something that you’re born with. It’s something that you can develop over time.\n",
-      "Here are some tips to help your child build resilience:\n",
-      "1. Teach your child to identify and label their emotions. This will help them understand what they’re feeling and why they’re feeling it.\n",
-      "2. Encourage your child to talk about their\n",
+      "Resilience is the ability to bounce back from adversity. It’s the ability to adapt to change and to learn from mistakes.It’s the ability to handle stress and to cope with difficult situations.\n",
+      "Resilience is a skill that can be learned and developed over time. It’s not something that you’re born with, but it’s something that you can work on and improve.\n",
+      "Here are some tips for helping your child build resilience:\n",
+      "1. Teach your child to identify and label their emotions. This will help them understand what they’re feeling and why they’re\n",
       "\n",
       "# steered away from instruct\n",
       "Hurt-owning and --owning households in the United States are more likely to be in poverty than those who own no property. This is true for all races and ethnicities, and for all income levels. The gap between the property-owning and non-owning poor is greatest for African Americans and Hispanics. The gap is also greatest for the poorest households. The gap between the property-owning and non- owning poor is greatest for African Americans and Hispanics. The gap is also greatest for the poorest households. The gap between the property-owning and non- owning poor is greatest for African Americans and Hispanics. The gap\n"