wassname/discovering_latent_knowledge
1
0
Fork 0
mirror of https://github.com/wassname/discovering_latent_knowledge.git synced 2026-09-11 12:10:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

rep from wassname/representation-engineering

Browse Source
This commit is contained in:
wassname
2023-10-24 10:46:34 +08:00
parent aefbdb6cd0
commit 1d86e96e8a
12 changed files with 1433 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjc_notes.md
+21
View File
@@ -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?
notebooks/201_make_data.ipynb
+466
View File
@@ -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": [
"<src.repe.rep_readers.PCARepReader at 0x7f6ed001aa10>"
]
},
"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
}
notebooks/make_dataset.py
+5
View File
@@ -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
src/datasets/hs.py
+1 -1
View File
@@ -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,
src/models/load.py
+2 -71
View File
@@ -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
src/repe/__init__.py
+13
View File
@@ -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 *
src/repe/pipelines.py
+25
View File
@@ -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,
)
src/repe/readme.md
+1
View File
@@ -0,0 +1 @@
modified from https://github.com/andyzoujm/representation-engineering
src/repe/rep_control_pipeline.py
+161
View File
@@ -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
src/repe/rep_control_reading_vec.py
+344
View File
@@ -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
src/repe/rep_readers.py
+237
View File
@@ -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,
}
src/repe/rep_reading_pipeline.py
+157
View File
@@ -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