From 1d86e96e8af999d34cbf881ed6f23ea7bf640290 Mon Sep 17 00:00:00 2001 From: wassname Date: Tue, 24 Oct 2023 10:46:34 +0800 Subject: [PATCH] rep from wassname/representation-engineering --- mjc_notes.md | 21 ++ notebooks/201_make_data.ipynb | 466 ++++++++++++++++++++++++++++ notebooks/make_dataset.py | 5 + src/datasets/hs.py | 2 +- src/models/load.py | 73 +---- src/repe/__init__.py | 13 + src/repe/pipelines.py | 25 ++ src/repe/readme.md | 1 + src/repe/rep_control_pipeline.py | 161 ++++++++++ src/repe/rep_control_reading_vec.py | 344 ++++++++++++++++++++ src/repe/rep_readers.py | 237 ++++++++++++++ src/repe/rep_reading_pipeline.py | 157 ++++++++++ 12 files changed, 1433 insertions(+), 72 deletions(-) create mode 100644 notebooks/201_make_data.ipynb create mode 100644 src/repe/__init__.py create mode 100644 src/repe/pipelines.py create mode 100644 src/repe/readme.md create mode 100644 src/repe/rep_control_pipeline.py create mode 100644 src/repe/rep_control_reading_vec.py create mode 100644 src/repe/rep_readers.py create mode 100644 src/repe/rep_reading_pipeline.py diff --git a/mjc_notes.md b/mjc_notes.md index b63ef7d..4809c93 100644 --- a/mjc_notes.md +++ b/mjc_notes.md @@ -1801,3 +1801,24 @@ After reading https://github.dev/andyzoujm/representation-engineering/tree/main/ - use the diff of hidden states - use their intervention pipeline - and chain to a dataset? + + +I probobly don't need to wrap the model... just use baukit instead? + +Ok what should my pipeline do? +- take in text (they all do) +- params + - no cache + - return hidden states +- subclass + - text_classification? maybe + - feature extr (no text) but simple + - text_gen: very complex... adds whole prompt +- output text, scores, choice_scores, and hidden states + - but text and scores for only for the last, generated token +- args: + - call: choice_ids, activations + + +TODO: +- [ ] change layer -1 to actual name? diff --git a/notebooks/201_make_data.ipynb b/notebooks/201_make_data.ipynb new file mode 100644 index 0000000..946f626 --- /dev/null +++ b/notebooks/201_make_data.ipynb @@ -0,0 +1,466 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use pipelines as this https://github.com/wassname/representation-engineering/blob/random_comments_ignore/examples/honesty/honesty.ipynb\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + } + ], + "source": [ + "# import your package\n", + "%load_ext autoreload\n", + "%autoreload 2\n", + "\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "from matplotlib import pyplot as plt\n", + "plt.style.use('ggplot')\n", + "\n", + "import os\n", + "from pathlib import Path\n", + "from tqdm.auto import tqdm\n", + "from loguru import logger\n", + "logger.add(os.sys.stderr, format=\"{time} {level} {message}\", level=\"INFO\")\n", + "\n", + "from typing import Optional, List, Dict, Union, Tuple, Callable, Iterable\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "from torch import Tensor\n", + "from torch import optim\n", + "from torch.utils.data import random_split, DataLoader, TensorDataset\n", + "\n", + "import transformers\n", + "from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM\n", + "from src.repe import repe_pipeline_registry\n", + "repe_pipeline_registry()\n", + "\n", + "from src.models.load import load_model\n", + "from src.extraction.config import ExtractConfig\n", + "from make_dataset import create_hs_ds, load_preproc_dataset\n", + "\n", + "# from sklearn.linear_model import LogisticRegression\n", + "# from sklearn.metrics import f1_score, roc_auc_score, accuracy_score\n", + "# from sklearn.preprocessing import RobustScaler\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "ExtractConfig(datasets=('amazon_polarity', 'super_glue:boolq', 'glue:qnli', 'imdb'), model='TheBloke/WizardCoder-Python-13B-V1.0-GPTQ', data_dirs=(), max_examples=(100, 100), num_shots=1, num_variants=-1, layers=(), seed=42, template_path=None, max_length=555)\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\u001b[32m2023-10-24 10:45:58.826\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging pad_token_id from 32000 to 0\u001b[0m\n", + "2023-10-24T10:45:58.826830+0800 INFO changing pad_token_id from 32000 to 0\n", + "\u001b[32m2023-10-24 10:45:58.827\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging padding_side from right to left\u001b[0m\n", + "2023-10-24T10:45:58.827930+0800 INFO changing padding_side from right to left\n", + "\u001b[32m2023-10-24 10:45:58.828\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging truncation_side from right to left\u001b[0m\n", + "2023-10-24T10:45:58.828427+0800 INFO changing truncation_side from right to left\n" + ] + } + ], + "source": [ + "# model_name_or_path = \"TheBloke/Wizard-Vicuna-30B-Uncensored-GPTQ\"\n", + "# model_name_or_path = \"TheBloke/Mistral-7B-Instruct-v0.1-GPTQ\"\n", + "model_name_or_path = \"TheBloke/WizardCoder-Python-13B-V1.0-GPTQ\"\n", + "\n", + "cfg = ExtractConfig(max_examples=(100, 100), model=model_name_or_path)\n", + "print(cfg)\n", + "\n", + "model, tokenizer = load_model(model_name_or_path)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "rep_token = -1\n", + "batch_size = 2\n", + "# hidden_layers = list(range(-1, -model.config.num_hidden_layers, -1))\n", + "# hidden_layers = [f\"model.layers.{i}\" for i in range(8, model.config.num_hidden_layers, 3)]\n", + "hidden_layers = list(range(8, model.config.num_hidden_layers, 3))\n", + "hidden_layers \n", + "\n", + "n_difference = 1\n", + "direction_method = 'pca'\n", + "rep_reading_pipeline = pipeline(\"rep-reading\", model=model, tokenizer=tokenizer)\n", + "rep_reading_pipeline\n", + "hidden_layers\n" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "tokenize: 100%|██████████| 302/302 [00:00<00:00, 3370.75 examples/s]\n", + "truncated: 100%|██████████| 302/302 [00:00<00:00, 3882.16 examples/s]\n", + "prompt_truncated: 100%|██████████| 302/302 [00:00<00:00, 551.44 examples/s]\n", + "choice_ids: 100%|██████████| 302/302 [00:00<00:00, 9920.74 examples/s]\n", + "Filter: 100%|██████████| 302/302 [00:00<00:00, 3779.73 examples/s]" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "removed truncated rows to leave: num_rows 97\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + }, + { + "data": { + "text/plain": [ + "Dataset({\n", + " features: ['ds_string', 'example_i', 'answer', 'question', 'answer_choices', 'template_name', 'label_true', 'label_instructed', 'instructed_to_lie', 'sys_instr_name', 'input_ids', 'attention_mask', 'truncated', 'prompt_truncated', 'choice_ids'],\n", + " num_rows: 97\n", + "})" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# load dataset\n", + "ds_name = 'imdb'\n", + "ds_tokens = load_preproc_dataset(ds_name, cfg, tokenizer)\n", + "ds_tokens\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Dataset({\n", + " features: ['ds_string', 'example_i', 'answer', 'question', 'answer_choices', 'template_name', 'label_true', 'label_instructed', 'instructed_to_lie', 'sys_instr_name', 'input_ids', 'attention_mask', 'truncated', 'prompt_truncated', 'choice_ids'],\n", + " num_rows: 54\n", + "})" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "N_fit_examples = 10\n", + "N_train_split = (len(ds_tokens) - N_fit_examples) //2\n", + "\n", + "# split the dataset, it's preshuffled\n", + "dataset_fit = ds_tokens.select(range(N_fit_examples))\n", + "dataset_train = ds_tokens.select(range(N_fit_examples, N_train_split))\n", + "dataset_test = ds_tokens.select(range(N_train_split, len(ds_tokens)))\n", + "dataset_test\n" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer_args=dict(padding=\"max_length\", max_length=cfg.max_length, truncation=True, add_special_tokens=True)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# fit\n", + "train_labels = dataset_fit['label_true']\n", + "honesty_rep_reader = rep_reading_pipeline.get_directions(\n", + " dataset_fit['question'], \n", + " rep_token=rep_token, \n", + " hidden_layers=hidden_layers, \n", + " n_difference=n_difference, \n", + " train_labels=dataset_fit['label_true'], \n", + " direction_method=direction_method,\n", + " batch_size=batch_size,\n", + " **tokenizer_args\n", + ")\n", + "honesty_rep_reader\n" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "# read direction for each example, layer\n", + "H_tests = rep_reading_pipeline(\n", + " dataset_train['question'], \n", + " rep_token=rep_token, \n", + " hidden_layers=hidden_layers, \n", + " rep_reader=honesty_rep_reader,\n", + " batch_size=batch_size, **tokenizer_args)\n", + "H_tests[0] # {Batch, layers}\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Control" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "# layer_id = hidden_layers\n", + "# block_name=\"decoder_block\"\n", + "# control_method=\"reading_vec\"\n", + "\n", + "# rep_control_pipeline = pipeline(\n", + "# \"rep-control\", \n", + "# model=model, \n", + "# tokenizer=tokenizer, \n", + "# layers=layer_id, max_length=cfg.max_length,\n", + "# control_method=control_method)\n", + "# rep_control_pipeline\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "# from re import S\n", + "\n", + "\n", + "# inputs = dataset_train[:2]\n", + "# coeff=8.0\n", + "# max_new_tokens=3\n", + "# text_gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens, use_cache=False, \n", + "# output_hidden_states=True, return_dict=True,\n", + "# )\n", + "\n", + "# activations = {}\n", + "# for layer in layer_id:\n", + "# activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer]).to(model.device).half()\n", + " \n", + "\n", + "# activations_neg = {k:-v for k,v in activations.items()}\n", + "\n", + "# model.eval()\n", + "# with torch.no_grad():\n", + "# baseline_outputs = rep_control_pipeline(inputs, batch_size=batch_size, **text_gen_kwargs)\n", + "# control_outputs = rep_control_pipeline(inputs, activations=activations, batch_size=batch_size, **text_gen_kwargs)\n", + "# control_outputs_neg = rep_control_pipeline(inputs, activations=activations_neg, batch_size=batch_size, **text_gen_kwargs)\n", + "\n", + "# for i,s,p,n in zip(inputs, baseline_outputs['text_ans'], control_outputs['text_ans'], control_outputs_neg['text_ans']):\n", + "# print(\"===== No Control =====\")\n", + "# print(S)\n", + "# print(f\"===== + Honesty Control =====\")\n", + "# print(p)\n", + "# print()\n", + "# print(f\"===== - Honesty Control =====\")\n", + "# print(n)\n", + "# print()\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# control v2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "layer_id = hidden_layers\n", + "\n", + "rep_control_pipeline = pipeline(\n", + " \"rep-control2\", \n", + " model=model, \n", + " tokenizer=tokenizer, \n", + " layers=layer_id, \n", + " max_length=cfg.max_length,)\n", + "rep_control_pipeline\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "inputs = dataset_train[:2]\n", + "coeff=8.0\n", + "max_new_tokens=1\n", + "text_gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens, use_cache=False, \n", + " output_hidden_states=True, return_dict=True, max_length=cfg.max_length,\n", + " )\n", + "\n", + "activations = {}\n", + "for layer in layer_id:\n", + " activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer]).to(model.device).half()\n", + " \n", + "\n", + "activations_neg = {k:-v for k,v in activations.items()}\n", + "\n", + "model.eval()\n", + "with torch.no_grad():\n", + " baseline_outputs = rep_control_pipeline(inputs, batch_size=batch_size, **text_gen_kwargs)\n", + " control_outputs = rep_control_pipeline(inputs, activations=activations, batch_size=batch_size, **text_gen_kwargs)\n", + " control_outputs_neg = rep_control_pipeline(inputs, activations=activations_neg, batch_size=batch_size, **text_gen_kwargs)\n", + "\n", + "\n", + "for i,s,p,n in zip(inputs, baseline_outputs, control_outputs, control_outputs_neg):\n", + " print(\"===== No Control =====\")\n", + " print(s['generated_text'][0].replace(i, \"\"))\n", + " print(f\"===== + Honesty Control =====\")\n", + " print(p['generated_text'][0].replace(i, \"\"))\n", + " print()\n", + " print(f\"===== - Honesty Control =====\")\n", + " print(n['generated_text'][0].replace(i, \"\"))\n", + " print()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "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.10.12" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/notebooks/make_dataset.py b/notebooks/make_dataset.py index 16748b7..45ec6ce 100644 --- a/notebooks/make_dataset.py +++ b/notebooks/make_dataset.py @@ -225,6 +225,8 @@ def qc_ds(f): def load_preproc_dataset(ds_name: str, cfg: ExtractConfig, tokenizer: PreTrainedTokenizerBase, split_type:str="train", N=None) -> Dataset: + """load a preprocessed dataset of tokens.""" + # TODO refactor out cfg if N is None: N = cfg.max_examples[split_type!="train"] ds_prompts = Dataset.from_generator( @@ -266,6 +268,8 @@ def load_preproc_dataset(ds_name: str, cfg: ExtractConfig, tokenizer: PreTrained .map(lambda r: {'choice_ids': row_choice_ids(r, tokenizer)}, desc='choice_ids') ) + + ds_tokens = shuffle_dataset_by(ds_tokens, 'example_i') @@ -323,6 +327,7 @@ def post_proc_hs_ds(ds1, tokenizer): return ds3 def create_hs_ds(ds_name, ds_tokens, model, cfg, intervention_dicts = [None, ], f = None, split_type="train"): + "create a dataset of hidden states.""" info_kwargs = dict(extract_cfg=cfg.to_dict(), ds_name=ds_name, split_type=split_type, f=f, date=pd.Timestamp.now().isoformat(),) # first we make the calibration dataset with no intervention diff --git a/src/datasets/hs.py b/src/datasets/hs.py index 4771fb9..7d7de76 100644 --- a/src/datasets/hs.py +++ b/src/datasets/hs.py @@ -148,7 +148,7 @@ class ExtractHiddenStates: with torch.no_grad(): multi_outs = defaultdict(list) for edit_output in edit_outputs: - with TraceDict(self.model, layers_names, retain_grad=True, detach=True, edit_output=edit_output) as ret: + with TraceDict(self.model, layers_names, retain_grad=False, detach=True, edit_output=edit_output) as ret: model_inputs = self.model.prepare_inputs_for_generation(input_ids=input_ids, attention_mask=attention_mask, use_cache=False) outputs = self.model.forward( **model_inputs, diff --git a/src/models/load.py b/src/models/load.py index 94dc03d..11cbeb8 100644 --- a/src/models/load.py +++ b/src/models/load.py @@ -21,94 +21,25 @@ def verbose_change_param(tokenizer, path, after): def load_model(model_repo = "TheBloke/WizardCoder-Python-13B-V1.0-GPTQ") -> Tuple[AutoModelForCausalLM, PreTrainedTokenizerBase]: """ - Chosing: - - https://old.reddit.com/r/LocalLLaMA/wiki/models - - https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard - - https://github.com/deep-diver/LLM-As-Chatbot/blob/main/model_cards.json - - A uncensored and large coding ones might be best for lying. """ # see https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/starchat.py # gptq_config = GPTQConfig(bits=4, dataset="c4", disable_exllama=False) model_options = dict( device_map="auto", - # load_in_8bit=True, - # load_in_4bit=True, - # torch_dtype=torch.float16, # note because datasets pickles the model into numpy to get the unique datasets name, and because numpy doesn't support bfloat16, we need to use float16 - # use_safetensors=True, + torch_dtype=torch.float16, ) config = AutoConfig.from_pretrained(model_repo, use_cache=False) verbose_change_param(config, 'use_cache', False) - tokenizer = AutoTokenizer.from_pretrained(model_repo, use_fast=True) + tokenizer = AutoTokenizer.from_pretrained(model_repo, use_fast=True, legacy=False) verbose_change_param(tokenizer, 'pad_token_id', 0) verbose_change_param(tokenizer, 'padding_side', 'left') verbose_change_param(tokenizer, 'truncation_side', 'left') model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, - # quantization_config = gptq_config, **model_options) return model, tokenizer - -# def load_model(model_repo = "HuggingFaceH4/starchat-beta", lora_repo=None, verbose=True): -# if "starchat" in model_repo: -# model, tokenizer = load_starchat(model_repo=model_repo) -# # elif "llama" in model_repo: -# # model, tokenizer = load_llama(model_repo=model_repo, lora_repo=lora_repo) -# else: -# raise NotImplementedError(f"code for model_repo {model_repo} not found") - -# if verbose: print(model.config) - -# assert check_for_dropout(model), 'model should have dropout' -# return model, tokenizer - -# def load_starchat(model_repo = "HuggingFaceH4/starchat-beta", load_in_4bit=True, torch_dtype=torch.float16): -# # see https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/starchat.py -# model_options = dict( -# device_map="auto", -# load_in_4bit=load_in_4bit, -# torch_dtype=torch_dtype, # note because datasets pickles the model into numpy to get the unique datasets name, and because numpy doesn't support bfloat16, we need to use float16 -# use_safetensors=False, -# ) - -# config = AutoConfig.from_pretrained(model_repo, use_cache=False) -# verbose_change_param(config, 'use_cache', False) - -# tokenizer = AutoTokenizer.from_pretrained(model_repo) -# verbose_change_param(tokenizer, 'pad_token_id', 0) -# verbose_change_param(tokenizer, 'padding_side', 'left') -# verbose_change_param(tokenizer, 'truncation_side', 'left') - -# model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options) - -# return model, tokenizer - -# def load_llama(model_repo, lora_repo=None): -# # https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/alpaca.py -# model_options = dict( -# device_map="auto", -# load_in_4bit=True, -# torch_dtype=torch.float16, -# ) - -# tokenizer = LlamaTokenizer.from_pretrained(model_repo) -# model = LlamaForCausalLM.from_pretrained(model_repo, **model_options) - -# if lora_repo is not None: -# # https://github.com/tloen/alpaca-lora/blob/main/generate.py#L40 -# from peft import PeftModel -# model = PeftModel.from_pretrained( -# model, -# lora_repo, -# torch_dtype=torch.float16, -# device_map='auto' -# ) -# return model, tokenizer - -# def load_falcan(model_repo, lora_repo=None): -# # https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/falcon.py diff --git a/src/repe/__init__.py b/src/repe/__init__.py new file mode 100644 index 0000000..faa7931 --- /dev/null +++ b/src/repe/__init__.py @@ -0,0 +1,13 @@ +import warnings +warnings.filterwarnings("ignore") + + +from .pipelines import repe_pipeline_registry + +# RepReading +from .rep_readers import * +from .rep_reading_pipeline import * + +# RepControl +from .rep_control_pipeline import * +from .rep_control_reading_vec import * \ No newline at end of file diff --git a/src/repe/pipelines.py b/src/repe/pipelines.py new file mode 100644 index 0000000..ababfbf --- /dev/null +++ b/src/repe/pipelines.py @@ -0,0 +1,25 @@ +from transformers import AutoModel, AutoModelForCausalLM +from transformers.pipelines import PIPELINE_REGISTRY +from .rep_reading_pipeline import RepReadingPipeline +from .rep_control_pipeline import RepControlPipeline, RepControlPipeline2 + +def repe_pipeline_registry(): + PIPELINE_REGISTRY.register_pipeline( + "rep-reading", + pipeline_class=RepReadingPipeline, + pt_model=AutoModel, + ) + + PIPELINE_REGISTRY.register_pipeline( + "rep-control", + pipeline_class=RepControlPipeline, + pt_model=AutoModelForCausalLM, + ) + + PIPELINE_REGISTRY.register_pipeline( + "rep-control2", + pipeline_class=RepControlPipeline2, + pt_model=AutoModelForCausalLM, + ) + + diff --git a/src/repe/readme.md b/src/repe/readme.md new file mode 100644 index 0000000..225deba --- /dev/null +++ b/src/repe/readme.md @@ -0,0 +1 @@ +modified from https://github.com/andyzoujm/representation-engineering diff --git a/src/repe/rep_control_pipeline.py b/src/repe/rep_control_pipeline.py new file mode 100644 index 0000000..e07f450 --- /dev/null +++ b/src/repe/rep_control_pipeline.py @@ -0,0 +1,161 @@ +import torch +from transformers.pipelines import ( + TextGenerationPipeline, + FeatureExtractionPipeline, + Pipeline, +) +from transformers.pipelines.base import GenericTensor +from .rep_control_reading_vec import WrappedReadingVecModel +from typing import Dict + + +class RepControlPipeline(FeatureExtractionPipeline): + def __init__( + self, + model, + tokenizer, + layers, + block_name="decoder_block", + control_method="reading_vec", + max_length=555, + **kwargs, + ): + # TODO: implement different control method and supported intermediate modules for different models + assert control_method == "reading_vec", f"{control_method} not supported yet" + assert ( + block_name == "decoder_block" + or "LlamaForCausalLM" in model.config.architectures + ), f"{model.config.architectures} {block_name} not supported yet" + self.wrapped_model = WrappedReadingVecModel(model, tokenizer) + self.wrapped_model.unwrap() + self.wrapped_model.wrap_block(layers, block_name=block_name) + self.block_name = block_name + self.layers = layers + self.max_length = max_length + + super().__init__(model=model, tokenizer=tokenizer, **kwargs) + + def preprocess(self, inputs, **tokenize_kwargs) -> Dict[str, GenericTensor]: + # tokenize a batch of inputs + return_tensors = self.framework + model_inputs = self.tokenizer(inputs['question'], return_tensors=return_tensors, return_attention_mask=True, add_special_tokens=True, truncation=True, padding="max_length", max_length=self.max_length, **tokenize_kwargs) + return {**inputs, **model_inputs} + + def __call__(self, text_inputs, activations=None, **kwargs): + if activations is not None: + self.wrapped_model.reset() + self.wrapped_model.set_controller(self.layers, activations, self.block_name) + + outputs = super().__call__(text_inputs, **kwargs) + self.wrapped_model.reset() + + return outputs + + def _forward(self, model_inputs): + inputs = dict(input_ids=model_inputs['input_ids'], attention_mask=model_inputs['attention_mask']) + inputs.update( + {"use_cache": False, "output_hidden_states": True, "return_dict": True} + ) + with torch.no_grad(): + model_outputs = self.model(**inputs) + + # retain some of the inputs + keep_cols = ["answer_choices", "input_ids", "attention_mask"] + model_outputs = {**model_inputs, **model_outputs} + return model_outputs + + def postprocess(self, o): + # note this sometimes deals with a batch, sometimes with a single result. infuriating + assert isinstance(o, dict) and o['logits'].ndim==3, f"expected dict with logits of shape (batch, seq, vocab), got {o['logits'].shape}" + # This is called once for each result, but the text pipeline is set up to hande multiple... + # This is called once for each result, but the text pipeline is set up to hande multiple... + o["end_logits"] = o["logits"][:, -1, :].float() + # hidden_states = list(o.hidden_states) + o["input_truncated"] = self.tokenizer.batch_decode(o['input_ids']) + o["truncated"] = torch.sum(o["attention_mask"], 1)==self.max_length + o["text_ans"] = self.tokenizer.batch_decode(o["end_logits"].argmax(-1)) + o['choice_ids'] = row_choice_ids(o, self.tokenizer) + return o + + +from typing import List, Tuple, Dict, Any, Union, NewType +from baukit.nethook import Trace, TraceDict, recursive_copy +from src.datasets.intervene import InterventionDict, intervention_meta_fn +from functools import partial +from src.datasets.scores import choice2ids + +Activations = NewType("InterventionDict", Dict[str, torch.Tensor]) + +def row_choice_ids(answer_choices, tokenizer): + return choice2ids([[c] for c in answer_choices], tokenizer) + +def intervention_meta_fn2( + output: torch.Tensor, layer_name: str, activations: Activations +) -> torch.Tensor: + """see + - honest_llama: https://github.com/likenneth/honest_llama/blob/e010f82bfbeaa4326cef8493b0dd5b8b14c6da67/validation/validate_2fold.py#L114 + - baukit: https://github.com/davidbau/baukit/blob/main/baukit/nethook.py#L42C1-L45C56 + + Usage: + edit_output = partial(intervention_meta_fn2, activations=activations) + with TraceDict(model, layers_to_intervene, edit_output=edit_output) as ret: + ... + + """ + for activation in activations[layer_name]: + # TODO might be model specific? + output[:, :, :] += torch.from_numpy(activation).to(output.device)[None, None, :] + + return output + + +class RepControlPipeline2(FeatureExtractionPipeline): + """This version uses baukit.""" + def __init__(self, model, tokenizer, max_length, **kwargs): + super().__init__(model=model, tokenizer=tokenizer, **kwargs) + self.max_length = max_length + + def __call__(self, model_inputs, activations=None, **kwargs): + if activations is not None: + # FIXME model specific + layers_names = [f'model.model.layers.{i}.post_attention_layernorm' for i in activations.keys()] + edit_fn = partial(intervention_meta_fn2, activations=activations) + with TraceDict( + self.model, layers_names, detach=True, edit_output=edit_fn + ) as ret: + outputs = super().__call__(model_inputs, **kwargs) + else: + outputs = super().__call__(model_inputs, **kwargs) + return outputs + + def preprocess(self, inputs, **tokenize_kwargs) -> Dict[str, GenericTensor]: + # tokenize a batch of inputs + return_tensors = self.framework + model_inputs = self.tokenizer(inputs['question'], return_tensors=return_tensors, return_attention_mask=True, add_special_tokens=True, truncation=True, padding="max_length", max_length=self.max_length, **tokenize_kwargs) + return {**inputs, **model_inputs} + + def _forward(self, model_inputs): + inputs = dict(input_ids=model_inputs['input_ids'], attention_mask=model_inputs['attention_mask']) + inputs.update( + {"use_cache": False, "output_hidden_states": True, "return_dict": True} + ) + with torch.no_grad(): + model_outputs = self.model(**inputs) + + # retain some of the inputs + keep_cols = ["answer_choices", "input_ids", "attention_mask"] + model_outputs = {**model_inputs, **model_outputs} + return model_outputs + + def postprocess(self, o): + # note this sometimes deals with a batch, sometimes with a single result. infuriating + assert isinstance(o, dict) and o['logits'].ndim==3, f"expected dict with logits of shape (batch, seq, vocab), got {o['logits'].shape}" + # This is called once for each result, but the text pipeline is set up to hande multiple... + # This is called once for each result, but the text pipeline is set up to hande multiple... + o["end_logits"] = o["logits"][:, -1, :].float() + # hidden_states = list(o.hidden_states) + o["input_truncated"] = self.tokenizer.batch_decode(o['input_ids']) + o["truncated"] = torch.sum(o["attention_mask"], 1)==self.max_length + o["text_ans"] = self.tokenizer.batch_decode(o["end_logits"].argmax(-1)) + o['choice_ids'] = row_choice_ids(o, self.tokenizer) + return o diff --git a/src/repe/rep_control_reading_vec.py b/src/repe/rep_control_reading_vec.py new file mode 100644 index 0000000..7d07cb8 --- /dev/null +++ b/src/repe/rep_control_reading_vec.py @@ -0,0 +1,344 @@ +# wrapping classes +import torch +import numpy as np + +class WrappedBlock(torch.nn.Module): + def __init__(self, block): + super().__init__() + self.block = block + self.output = None + self.controller = None + self.mask = None + self.token_pos = None + self.normalize = False + + def forward(self, *args, **kwargs): + output = self.block(*args, **kwargs) + + if isinstance(output, tuple): + self.output = output[0] + modified = output[0] + else: + self.output = output + modified = output + + + if self.controller is not None: + + norm_pre = torch.norm(modified, dim=-1, keepdim=True) + + if self.mask is not None: + mask = self.mask + + # we should ignore the padding tokens when doing the activation addition + # mask has ones for non padding tokens and zeros at padding tokens. + # only tested this on left padding + elif "position_ids" in kwargs: + pos = kwargs["position_ids"] + zero_indices = (pos == 0).cumsum(1).argmax(1, keepdim=True) + col_indices = torch.arange(pos.size(1), device=pos.device).unsqueeze(0) + target_shape = pos.shape + mask = (col_indices >= zero_indices).float().reshape(target_shape[0], target_shape[1], 1) + mask = mask.to(modified.dtype) + else: + # print(f"Warning: block {self.block_name} does not contain information 'position_ids' about token types. When using batches this can lead to unexpected results.") + mask = 1.0 + + if len(self.controller.shape) == 1: + self.controller = self.controller.reshape(1, 1, -1) + assert len(self.controller.shape) == len(modified.shape), f"Shape of controller {self.controller.shape} does not match shape of modified {modified.shape}." + + self.controller = self.controller.to(modified.device) + if type(mask) == torch.Tensor: + mask = mask.to(modified.device) + if isinstance(self.token_pos, int): + modified[:, self.token_pos] = modified[:, self.token_pos] + self.controller * mask + elif isinstance(self.token_pos, list) or isinstance(self.token_pos, tuple) or isinstance(self.token_pos, np.ndarray): + modified[:, self.token_pos] = modified[:, self.token_pos] + self.controller * mask + elif isinstance(self.token_pos, str): + if self.token_pos == "end": + len_token = self.controller.shape[1] + modified[:, -len_token:] = modified[:, -len_token:] + self.controller * mask + elif self.token_pos == "start": + len_token = self.controller.shape[1] + modified[:, :len_token] = modified[:, :len_token] + self.controller * mask + else: + assert False, f"Unknown token position {self.token_pos}." + else: + modified = modified + self.controller * mask + + if self.normalize: + norm_post = torch.norm(modified, dim=-1, keepdim=True) + modified = modified / norm_post * norm_pre + + if isinstance(output, tuple): + output = (modified,) + output[1:] + else: + output = modified + + return output + + def set_controller(self, activations, token_pos=None, masks=None, normalize=False): + self.normalize = normalize + self.controller = activations.squeeze() + self.mask = masks + self.token_pos = token_pos + + def reset(self): + self.output = None + self.controller = None + self.mask = None + + def set_masks(self, masks): + self.mask = masks + + +class WrappedReadingVecModel(torch.nn.Module): + def __init__(self, model, tokenizer): + super().__init__() + self.model = model + self.tokenizer = tokenizer + + def forward(self, *args, **kwargs): + return self.model(*args, **kwargs) + + def generate(self, prompt, max_new_tokens=100, random_seed=0, use_cache=True): + with torch.no_grad(): + torch.random.manual_seed(random_seed) + inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, max_length=512, truncation=True) + attention_mask = inputs.attention_mask.to(self.model.device) + generate_ids = self.model.generate(inputs.input_ids.to(self.model.device), attention_mask=attention_mask, max_new_tokens=max_new_tokens, use_cache=use_cache) + return self.tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) + + def get_logits(self, tokens): + with torch.no_grad(): + logits = self.model(tokens.to(self.model.device)).logits + return logits + + def run_prompt(self, prompt, **kwargs): + with torch.no_grad(): + inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, max_length=512, truncation=True) + input_ids = inputs.input_ids.to(self.model.device) + attention_mask = inputs.attention_mask.to(self.model.device) + output = self.model(input_ids, attention_mask=attention_mask) + return output + + def wrap_self_attn(self, layer_id): + if self.is_wrapped(self.model.model.layers[layer_id]): + block = self.model.model.layers[layer_id].block.self_attn + if not self.is_wrapped(block): + self.model.model.layers[layer_id].block.self_attn = WrappedBlock(block) + else: + block = self.model.model.layers[layer_id].self_attn + if not self.is_wrapped(block): + self.model.model.layers[layer_id].self_attn = WrappedBlock(block) + + def wrap_mlp(self, layer_id): + if self.is_wrapped(self.model.model.layers[layer_id]): + block = self.model.model.layers[layer_id].block.mlp + if not self.is_wrapped(block): + self.model.model.layers[layer_id].block.mlp = WrappedBlock(block) + else: + block = self.model.model.layers[layer_id].mlp + if not self.is_wrapped(block): + self.model.model.layers[layer_id].mlp = WrappedBlock(block) + + def wrap_input_layernorm(self, layer_id): + if self.is_wrapped(self.model.model.layers[layer_id]): + block = self.model.model.layers[layer_id].block.input_layernorm + if not self.is_wrapped(block): + self.model.model.layers[layer_id].block.input_layernorm = WrappedBlock(block) + else: + block = self.model.model.layers[layer_id].input_layernorm + if not self.is_wrapped(block): + self.model.model.layers[layer_id].input_layernorm = WrappedBlock(block) + + def wrap_post_attention_layernorm(self, layer_id): + if self.is_wrapped(self.model.model.layers[layer_id]): + block = self.model.model.layers[layer_id].block.post_attention_layernorm + if not self.is_wrapped(block): + self.model.model.layers[layer_id].block.post_attention_layernorm = WrappedBlock(block) + else: + block = self.model.model.layers[layer_id].post_attention_layernorm + if not self.is_wrapped(block): + self.model.model.layers[layer_id].post_attention_layernorm = WrappedBlock(block) + + def wrap_decoder_block(self, layer_id): + block = self.model.model.layers[layer_id] + if not self.is_wrapped(block): + self.model.model.layers[layer_id] = WrappedBlock(block) + + + def wrap_all(self): + for layer_id, layer in enumerate(self.model.model.layers): + self.wrap_self_attn(layer_id) + self.wrap_mlp(layer_id) + self.wrap_input_layernorm(layer_id) + self.wrap_post_attention_layernorm(layer_id) + self.wrap_decoder_block(layer_id) + + def wrap_block(self, layer_ids, block_name): + def _wrap_block(layer_id, block_name): + if block_name == 'self_attn': + self.wrap_self_attn(layer_id) + elif block_name == 'mlp': + self.wrap_mlp(layer_id) + elif block_name == 'input_layernorm': + self.wrap_input_layernorm(layer_id) + elif block_name == 'post_attention_layernorm': + self.wrap_post_attention_layernorm(layer_id) + elif block_name == 'decoder_block': + self.wrap_decoder_block(layer_id) + else: + assert False, f"No block named {block_name}." + + if isinstance(layer_ids, list) or isinstance(layer_ids, tuple) or isinstance(layer_ids, np.ndarray): + for layer_id in layer_ids: + _wrap_block(layer_id, block_name) + else: + _wrap_block(layer_ids, block_name) + + + def get_activations(self, layer_ids, block_name='decoder_block'): + + def _get_activations(layer_id, block_name): + current_layer = self.model.model.layers[layer_id] + + if self.is_wrapped(current_layer): + current_block = current_layer.block + if block_name == 'decoder_block': + return current_layer.output + elif block_name == 'self_attn' and self.is_wrapped(current_block.self_attn): + return current_block.self_attn.output + elif block_name == 'mlp' and self.is_wrapped(current_block.mlp): + return current_block.mlp.output + elif block_name == 'input_layernorm' and self.is_wrapped(current_block.input_layernorm): + return current_block.input_layernorm.output + elif block_name == 'post_attention_layernorm' and self.is_wrapped(current_block.post_attention_layernorm): + return current_block.post_attention_layernorm.output + else: + assert False, f"No wrapped block named {block_name}." + + else: + if block_name == 'self_attn' and self.is_wrapped(current_layer.self_attn): + return current_layer.self_attn.output + elif block_name == 'mlp' and self.is_wrapped(current_layer.mlp): + return current_layer.mlp.output + elif block_name == 'input_layernorm' and self.is_wrapped(current_layer.input_layernorm): + return current_layer.input_layernorm.output + elif block_name == 'post_attention_layernorm' and self.is_wrapped(current_layer.post_attention_layernorm): + return current_layer.post_attention_layernorm.output + else: + assert False, f"No wrapped block named {block_name}." + + if isinstance(layer_ids, list) or isinstance(layer_ids, tuple) or isinstance(layer_ids, np.ndarray): + activations = {} + for layer_id in layer_ids: + activations[layer_id] = _get_activations(layer_id, block_name) + return activations + else: + return _get_activations(layer_ids, block_name) + + + def set_controller(self, layer_ids, activations, block_name='decoder_block', token_pos=None, masks=None, normalize=False): + + def _set_controller(layer_id, activations, block_name, masks, normalize): + current_layer = self.model.model.layers[layer_id] + + if block_name == 'decoder_block': + current_layer.set_controller(activations, token_pos, masks, normalize) + elif self.is_wrapped(current_layer): + current_block = current_layer.block + if block_name == 'self_attn' and self.is_wrapped(current_block.self_attn): + current_block.self_attn.set_controller(activations, token_pos, masks, normalize) + elif block_name == 'mlp' and self.is_wrapped(current_block.mlp): + current_block.mlp.set_controller(activations, token_pos, masks, normalize) + elif block_name == 'input_layernorm' and self.is_wrapped(current_block.input_layernorm): + current_block.input_layernorm.set_controller(activations, token_pos, masks, normalize) + elif block_name == 'post_attention_layernorm' and self.is_wrapped(current_block.post_attention_layernorm): + current_block.post_attention_layernorm.set_controller(activations, token_pos, masks, normalize) + else: + return f"No wrapped block named {block_name}." + + else: + if block_name == 'self_attn' and self.is_wrapped(current_layer.self_attn): + current_layer.self_attn.set_controller(activations, token_pos, masks, normalize) + elif block_name == 'mlp' and self.is_wrapped(current_layer.mlp): + current_layer.mlp.set_controller(activations, token_pos, masks, normalize) + elif block_name == 'input_layernorm' and self.is_wrapped(current_layer.input_layernorm): + current_layer.input_layernorm.set_controller(activations, token_pos, masks, normalize) + elif block_name == 'post_attention_layernorm' and self.is_wrapped(current_layer.post_attention_layernorm): + current_layer.post_attention_layernorm.set_controller(activations, token_pos, masks, normalize) + else: + return f"No wrapped block named {block_name}." + + if isinstance(layer_ids, list) or isinstance(layer_ids, tuple) or isinstance(layer_ids, np.ndarray): + assert isinstance(activations, dict), "activations should be a dictionary" + for layer_id in layer_ids: + _set_controller(layer_id, activations[layer_id], block_name, masks, normalize) + else: + _set_controller(layer_ids, activations, block_name, masks, normalize) + + + def reset(self): + for layer in self.model.model.layers: + if self.is_wrapped(layer): + layer.reset() + if self.is_wrapped(layer.block.self_attn): + layer.block.self_attn.reset() + if self.is_wrapped(layer.block.mlp): + layer.block.mlp.reset() + if self.is_wrapped(layer.block.input_layernorm): + layer.block.input_layernorm.reset() + if self.is_wrapped(layer.block.post_attention_layernorm): + layer.block.post_attention_layernorm.reset() + else: + if self.is_wrapped(layer.self_attn): + layer.self_attn.reset() + if self.is_wrapped(layer.mlp): + layer.mlp.reset() + if self.is_wrapped(layer.input_layernorm): + layer.input_layernorm.reset() + if self.is_wrapped(layer.post_attention_layernorm): + layer.post_attention_layernorm.reset() + + def set_masks(self, masks): + for layer in self.model.model.layers: + if self.is_wrapped(layer): + layer.set_masks(masks) + if self.is_wrapped(layer.block.self_attn): + layer.block.self_attn.set_masks(masks) + if self.is_wrapped(layer.block.mlp): + layer.block.mlp.set_masks(masks) + if self.is_wrapped(layer.block.input_layernorm): + layer.block.input_layernorm.set_masks(masks) + if self.is_wrapped(layer.block.post_attention_layernorm): + layer.block.post_attention_layernorm.set_masks(masks) + else: + if self.is_wrapped(layer.self_attn): + layer.self_attn.set_masks(masks) + if self.is_wrapped(layer.mlp): + layer.mlp.set_masks(masks) + if self.is_wrapped(layer.input_layernorm): + layer.input_layernorm.set_masks(masks) + if self.is_wrapped(layer.post_attention_layernorm): + layer.post_attention_layernorm.set_masks(masks) + + + def is_wrapped(self, block): + if hasattr(block, 'block'): + return True + return False + + def unwrap(self): + for l, layer in enumerate(self.model.model.layers): + if self.is_wrapped(layer): + self.model.model.layers[l] = layer.block + if self.is_wrapped(self.model.model.layers[l].self_attn): + self.model.model.layers[l].self_attn = self.model.model.layers[l].self_attn.block + if self.is_wrapped(self.model.model.layers[l].mlp): + self.model.model.layers[l].mlp = self.model.model.layers[l].mlp.block + if self.is_wrapped(self.model.model.layers[l].input_layernorm): + self.model.model.layers[l].input_layernorm = self.model.model.layers[l].input_layernorm.block + if self.is_wrapped(self.model.model.layers[l].post_attention_layernorm): + self.model.model.layers[l].post_attention_layernorm = self.model.model.layers[l].post_attention_layernorm.block diff --git a/src/repe/rep_readers.py b/src/repe/rep_readers.py new file mode 100644 index 0000000..0617d57 --- /dev/null +++ b/src/repe/rep_readers.py @@ -0,0 +1,237 @@ +from abc import ABC, abstractmethod +from sklearn.decomposition import PCA +from sklearn.cluster import KMeans +import numpy as np +from itertools import islice + +### Util Functions ### +def project_onto_direction(H, direction): + """Project matrix H (n, d_1) onto direction vector (d_2,)""" + # TODO: should we require direction vectors to be unit vectors? then return H.dot(direction) + mag = np.linalg.norm(direction) + assert not np.isinf(mag) + return H.dot(direction) / mag + +def recenter_mean(x, mean=None): + if mean is None: + mean = x.mean(axis=0, keepdims=True) + return x - mean + +class RepReader(ABC): + """Class to identify and store concept directions. + + Subclasses implement the abstract methods to identify concept directions + for each hidden layer via strategies including PCA, embedding vectors + (aka the logits method), and cluster means. + + RepReader instances are used by RepReaderPipeline to get concept scores. + + Directions can be used for downstream interventions.""" + + @abstractmethod + def __init__(self) -> None: + self.direction_method = None + self.directions = None # directions accessible via directions[layer][component_index] + self.direction_signs = None # direction of high concept scores (mapping min/max to high/low) + + @abstractmethod + def get_rep_directions(self, model, tokenizer, hidden_states, hidden_layers, **kwargs): + """Get concept directions for each hidden layer of the model + + Args: + model: Model to get directions for + tokenizer: Tokenizer to use + hidden_states: Hidden states of the model on the training data (per layer) + hidden_layers: Layers to consider + + Returns: + directions: A dict mapping layers to direction arrays (n_components, hidden_size) + """ + pass + + def get_signs(self, hidden_states, train_choices, hidden_layers): + """Given labels for the training data hidden_states, determine whether the + negative or positive direction corresponds to low/high concept + (and return corresponding signs -1 or 1 for each layer and component index) + + NOTE: This method assumes that there are 2 entries in hidden_states per label, + aka len(hidden_states[layer]) == 2 * len(train_choices). For example, if + n_difference=1, then hidden_states here should be the raw hidden states + rather than the relative (i.e. the differences between pairs of examples). + + Args: + hidden_states: Hidden states of the model on the training data (per layer) + train_choices: Labels for the training data + hidden_layers: Layers to consider + + Returns: + signs: A dict mapping layers to sign arrays (n_components,) + """ + signs = {} + + if self.needs_hiddens and hidden_states is not None and len(hidden_states) > 0: + for layer in hidden_layers: + assert hidden_states[layer].shape[0] == 2 * len(train_choices), f"Shape mismatch between hidden states ({hidden_states[layer].shape[0]}) and labels ({len(train_choices)})" + + signs[layer] = [] + for component_index in range(self.n_components): + transformed_hidden_states = project_onto_direction(hidden_states[layer], self.directions[layer][component_index]) + projected_scores = [transformed_hidden_states[i:i+2] for i in range(0, len(transformed_hidden_states), 2)] + + outputs_min = [1 if min(o) == o[label] else 0 for o, label in zip(projected_scores, train_choices)] + outputs_max = [1 if max(o) == o[label] else 0 for o, label in zip(projected_scores, train_choices)] + + signs[layer].append(-1 if np.mean(outputs_min) > np.mean(outputs_max) else 1) + else: + for layer in hidden_layers: + signs[layer] = [1 for _ in range(self.n_components)] + + return signs + + + def transform(self, hidden_states, hidden_layers, component_index): + """Project the hidden states onto the concept directions in self.directions + + Args: + hidden_states: dictionary with entries of dimension (n_examples, hidden_size) + hidden_layers: list of layers to consider + component_index: index of the component to use from self.directions + + Returns: + transformed_hidden_states: dictionary with entries of dimension (n_examples,) + """ + + assert component_index < self.n_components + + transformed_hidden_states = {} + for layer in hidden_layers: + layer_hidden_states = hidden_states[layer] + + if hasattr(self, 'H_train_means'): + layer_hidden_states = recenter_mean(layer_hidden_states, mean=self.H_train_means[layer]) + + # project hidden states onto found concept directions (e.g. onto PCA comp 0) + H_transformed = project_onto_direction(layer_hidden_states, self.directions[layer][component_index]) + transformed_hidden_states[layer] = H_transformed + + return transformed_hidden_states + +class PCARepReader(RepReader): + """Extract directions via PCA""" + needs_hiddens = True + + def __init__(self, n_components=1): + super().__init__() + self.n_components = n_components + self.H_train_means = {} + + def get_rep_directions(self, model, tokenizer, hidden_states, hidden_layers, **kwargs): + """Get PCA components for each layer""" + directions = {} + + for layer in hidden_layers: + H_train = np.array(hidden_states[layer]) + + H_train_mean = H_train.mean(axis=0, keepdims=True) + self.H_train_means[layer] = H_train_mean + H_train = recenter_mean(H_train, mean=H_train_mean) + + pca_model = PCA(n_components=self.n_components, whiten=False).fit(H_train) + + directions[layer] = pca_model.components_ # shape (n_components, n_features) + self.n_components = pca_model.n_components_ + + return directions + + def get_signs(self, hidden_states, train_labels, hidden_layers): + + signs = {} + + # WHY DO I NEED THIS FIXME? + # train_labels = np.array(train_labels)[:, None].tolist() + train_labels = [train_labels] + + for layer in hidden_layers: + assert hidden_states[layer].shape[0] == len(np.concatenate(train_labels)), f"Shape mismatch between hidden states ({hidden_states[layer].shape[0]}) and labels ({len(np.concatenate(train_labels))})" + layer_hidden_states = hidden_states[layer] + + # NOTE: since scoring is ultimately comparative, the effect of this is moot + layer_hidden_states = recenter_mean(layer_hidden_states, mean=self.H_train_means[layer]) + + # get the signs for each component + layer_signs = np.zeros(self.n_components) + for component_index in range(self.n_components): + + transformed_hidden_states = project_onto_direction(layer_hidden_states, self.directions[layer][component_index]) + + pca_outputs_comp = [list(islice(transformed_hidden_states, sum(len(c) for c in train_labels[:i]), sum(len(c) for c in train_labels[:i+1]))) for i in range(len(train_labels))] + + # We do elements instead of argmin/max because sometimes we pad random choices in training + pca_outputs_min = np.mean([o[train_labels[i].index(1)] == min(o) for i, o in enumerate(pca_outputs_comp)]) + pca_outputs_max = np.mean([o[train_labels[i].index(1)] == max(o) for i, o in enumerate(pca_outputs_comp)]) + + + layer_signs[component_index] = np.sign(np.mean(pca_outputs_max) - np.mean(pca_outputs_min)) + if layer_signs[component_index] == 0: + layer_signs[component_index] = 1 # default to positive in case of tie + + signs[layer] = layer_signs + + return signs + +class ClusterMeanRepReader(RepReader): + """Get the direction that is the difference between the mean of the positive and negative clusters.""" + n_components = 1 + needs_hiddens = True + + def __init__(self): + super().__init__() + + def get_rep_directions(self, model, tokenizer, hidden_states, hidden_layers, **kwargs): + + # train labels is necessary to differentiate between different classes + train_choices = kwargs['train_choices'] if 'train_choices' in kwargs else None + assert train_choices is not None, "ClusterMeanRepReader requires train_choices to differentiate two clusters" + for layer in hidden_layers: + assert len(train_choices) == len(hidden_states[layer]), f"Shape mismatch between hidden states ({len(hidden_states[layer])}) and labels ({len(train_choices)})" + + train_choices = np.array(train_choices) + neg_class = np.where(train_choices == 0) + pos_class = np.where(train_choices == 1) + + directions = {} + for layer in hidden_layers: + H_train = np.array(hidden_states[layer]) + + H_pos_mean = H_train[pos_class].mean(axis=0, keepdims=True) + H_neg_mean = H_train[neg_class].mean(axis=0, keepdims=True) + + directions[layer] = H_pos_mean - H_neg_mean + + return directions + + +class RandomRepReader(RepReader): + """Get random directions for each hidden layer. Do not use hidden + states or train labels of any kind.""" + + def __init__(self, needs_hiddens=True): + super().__init__() + + self.n_components = 1 + self.needs_hiddens = needs_hiddens + + def get_rep_directions(self, model, tokenizer, hidden_states, hidden_layers, **kwargs): + + directions = {} + for layer in hidden_layers: + directions[layer] = np.expand_dims(np.random.randn(model.config.hidden_size), 0) + + return directions + + +DIRECTION_FINDERS = { + 'pca': PCARepReader, + 'cluster_mean': ClusterMeanRepReader, + 'random': RandomRepReader, +} diff --git a/src/repe/rep_reading_pipeline.py b/src/repe/rep_reading_pipeline.py new file mode 100644 index 0000000..a71d137 --- /dev/null +++ b/src/repe/rep_reading_pipeline.py @@ -0,0 +1,157 @@ +from typing import List, Union, Optional +from transformers import Pipeline +import torch +import numpy as np +from .rep_readers import DIRECTION_FINDERS, RepReader + +class RepReadingPipeline(Pipeline): + """Returns the directions for each layer, for each example.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def _get_hidden_states( + self, + outputs, + rep_token: Union[str, int]=-1, + hidden_layers: Union[List[int], int]=-1, + which_hidden_states: Optional[str]=None): + + if hasattr(outputs, 'encoder_hidden_states') and hasattr(outputs, 'decoder_hidden_states'): + outputs['hidden_states'] = outputs[f'{which_hidden_states}_hidden_states'] + + hidden_states_layers = {} + for layer in hidden_layers: + hidden_states = outputs['hidden_states'][layer] + hidden_states = hidden_states[:, rep_token, :] + hidden_states_layers[layer] = hidden_states.cpu().to(dtype=torch.float32).detach().numpy() + + return hidden_states_layers + + def _sanitize_parameters(self, + rep_reader: RepReader=None, + rep_token: Union[str, int]=-1, + hidden_layers: Union[List[int], int]=-1, + component_index: int=0, + which_hidden_states: Optional[str]=None, + **tokenizer_kwargs): + preprocess_params = tokenizer_kwargs + forward_params = {} + postprocess_params = {} + + forward_params['rep_token'] = rep_token + + if not isinstance(hidden_layers, list): + hidden_layers = [hidden_layers] + + + assert rep_reader is None or len(rep_reader.directions) == len(hidden_layers), f"expect total rep_reader directions ({len(rep_reader.directions)})== total hidden_layers ({len(hidden_layers)})" + forward_params['rep_reader'] = rep_reader + forward_params['hidden_layers'] = hidden_layers + forward_params['component_index'] = component_index + forward_params['which_hidden_states'] = which_hidden_states + + return preprocess_params, forward_params, postprocess_params + + def preprocess( + self, + inputs: Union[str, List[str], List[List[str]]], + **tokenizer_kwargs): + + if self.image_processor: + return self.image_processor(inputs, add_end_of_utterance_token=False, return_tensors="pt") + return self.tokenizer(inputs, return_tensors=self.framework, **tokenizer_kwargs) + + def postprocess(self, outputs): + return outputs + + def _forward(self, model_inputs, rep_token, hidden_layers, rep_reader=None, component_index=0, which_hidden_states=None): + """ + Args: + - which_hidden_states (str): Specifies which part of the model (encoder, decoder, or both) to compute the hidden states from. + It's applicable only for encoder-decoder models. Valid values: 'encoder', 'decoder'. + """ + # get model hidden states and optionally transform them with a RepReader + with torch.no_grad(): + if hasattr(self.model, "encoder") and hasattr(self.model, "decoder"): + decoder_start_token = [self.tokenizer.pad_token] * model_inputs['input_ids'].size(0) + decoder_input = self.tokenizer(decoder_start_token, return_tensors="pt").input_ids + model_inputs['decoder_input_ids'] = decoder_input + outputs = self.model(**model_inputs, output_hidden_states=True) + hidden_states = self._get_hidden_states(outputs, rep_token, hidden_layers, which_hidden_states) + + if rep_reader is None: + return hidden_states + + return rep_reader.transform(hidden_states, hidden_layers, component_index) + + + def _batched_string_to_hiddens(self, train_inputs, rep_token, hidden_layers, batch_size, which_hidden_states, **tokenizer_args): + # Wrapper method to get a dictionary hidden states from a list of strings + hidden_states_outputs = self(train_inputs, rep_token=rep_token, + hidden_layers=hidden_layers, batch_size=batch_size, rep_reader=None, which_hidden_states=which_hidden_states, **tokenizer_args) + hidden_states = {layer: [] for layer in hidden_layers} + for hidden_states_batch in hidden_states_outputs: + for layer in hidden_states_batch: + hidden_states[layer].extend(hidden_states_batch[layer]) + return {k: np.array(v) for k, v in hidden_states.items()} + + def _validate_params(self, n_difference, direction_method): + # validate params for get_directions + if direction_method == 'clustermean': + assert n_difference == 1, "n_difference must be 1 for clustermean" + + def get_directions( + self, + train_inputs: Union[str, List[str], List[List[str]]], + rep_token: Union[str, int]=-1, + hidden_layers: Union[str, int]=-1, + n_difference: int = 1, + batch_size: int = 8, + train_labels: List[int] = None, + direction_method: str = 'pca', + direction_finder_kwargs: dict = {}, + which_hidden_states: Optional[str]=None, + **tokenizer_args,): + """Train a RepReader on the training data. + Args: + batch_size: batch size to use when getting hidden states + direction_method: string specifying the RepReader strategy for finding directions + direction_finder_kwargs: kwargs to pass to RepReader constructor + """ + + if not isinstance(hidden_layers, list): + assert isinstance(hidden_layers, int) + hidden_layers = [hidden_layers] + + self._validate_params(n_difference, direction_method) + + # initialize a DirectionFinder + direction_finder = DIRECTION_FINDERS[direction_method](**direction_finder_kwargs) + + # if relevant, get the hidden state data for training set + hidden_states = None + relative_hidden_states = None + if direction_finder.needs_hiddens: + # get raw hidden states for the train inputs + hidden_states = self._batched_string_to_hiddens(train_inputs, rep_token, hidden_layers, batch_size, which_hidden_states, **tokenizer_args) + + # get differences between pairs + relative_hidden_states = {k: np.copy(v) for k, v in hidden_states.items()} + for layer in hidden_layers: + for _ in range(n_difference): + relative_hidden_states[layer] = relative_hidden_states[layer][::2] - relative_hidden_states[layer][1::2] + + # get the directions + direction_finder.directions = direction_finder.get_rep_directions( + self.model, self.tokenizer, relative_hidden_states, hidden_layers, + train_choices=train_labels) + for layer in direction_finder.directions: + if type(direction_finder.directions[layer]) == np.ndarray: + direction_finder.directions[layer] = direction_finder.directions[layer].astype(np.float32) + + if train_labels is not None: + direction_finder.direction_signs = direction_finder.get_signs( + hidden_states, train_labels, hidden_layers) + + return direction_finder