Add hook for custom compute_hiddens

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
 

notebooks/model_delta.ipynb

@@ -0,0 +1,248 @@
+{
+ "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": 2,
+   "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": "8bbb21e1a2ae456aacdf88b9af4c94d6",
+       "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": "08e53b2b34a247aea7a5a42b9acf3b28",
+       "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": 5,
+   "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.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"
+     ]
+    }
+   ],
+   "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": 6,
+   "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. 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 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"
+     ]
+    }
+   ],
+   "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
+}

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,13 +244,25 @@ 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,
     inputs: list[DatasetEntry],
     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)

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