diff --git a/notebooks/201_make_data.ipynb b/notebooks/201_make_data.ipynb index 946f626..aa8f757 100644 --- a/notebooks/201_make_data.ipynb +++ b/notebooks/201_make_data.ipynb @@ -87,12 +87,12 @@ "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" + "\u001b[32m2023-10-24 10:49:20.938\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:49:20.938690+0800 INFO changing pad_token_id from 32000 to 0\n", + "\u001b[32m2023-10-24 10:49:20.940\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:49:20.940322+0800 INFO changing padding_side from right to left\n", + "\u001b[32m2023-10-24 10:49:20.941\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:49:20.941383+0800 INFO changing truncation_side from right to left\n" ] } ], @@ -147,11 +147,11 @@ "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]" + "tokenize: 100%|██████████| 302/302 [00:00<00:00, 3490.12 examples/s]\n", + "truncated: 100%|██████████| 302/302 [00:00<00:00, 3889.90 examples/s]\n", + "prompt_truncated: 100%|██████████| 302/302 [00:00<00:00, 604.08 examples/s]\n", + "choice_ids: 100%|██████████| 302/302 [00:00<00:00, 9749.77 examples/s]\n", + "Filter: 100%|██████████| 302/302 [00:00<00:00, 3987.01 examples/s]" ] }, { @@ -236,7 +236,7 @@ { "data": { "text/plain": [ - "" + "" ] }, "execution_count": 8, @@ -264,7 +264,28 @@ "cell_type": "code", "execution_count": 9, "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/plain": [ + "{8: -7.1495447,\n", + " 11: 16.796896,\n", + " 14: 13.5725565,\n", + " 17: -5.684828,\n", + " 20: 26.571663,\n", + " 23: 37.01544,\n", + " 26: -67.76076,\n", + " 29: -64.502205,\n", + " 32: -52.72602,\n", + " 35: -76.70331,\n", + " 38: -80.74953}" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "# read direction for each example, layer\n", "H_tests = rep_reading_pipeline(\n", @@ -285,14 +306,7 @@ }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, + "execution_count": 10, "metadata": {}, "outputs": [], "source": [ @@ -312,14 +326,7 @@ }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, + "execution_count": 11, "metadata": {}, "outputs": [], "source": [ @@ -329,7 +336,7 @@ "\n", "# inputs = dataset_train[:2]\n", "# coeff=8.0\n", - "# max_new_tokens=3\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,\n", "# )\n", @@ -349,7 +356,7 @@ "\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(s)\n", "# print(f\"===== + Honesty Control =====\")\n", "# print(p)\n", "# print()\n", @@ -374,9 +381,20 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 12, "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "layer_id = hidden_layers\n", "\n", @@ -391,13 +409,36 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 47, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "===== No Control =====\n", + "negative\n", + "===== + Honesty Control =====\n", + "negative\n", + "\n", + "===== - Honesty Control =====\n", + "negative\n", + "\n", + "===== No Control =====\n", + "The\n", + "===== + Honesty Control =====\n", + "The\n", + "\n", + "===== - Honesty Control =====\n", + "negative\n", + "\n" + ] + } + ], "source": [ "inputs = dataset_train[:2]\n", "coeff=8.0\n", - "max_new_tokens=1\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, max_length=cfg.max_length,\n", " )\n", @@ -416,23 +457,56 @@ " 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", + "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['generated_text'][0].replace(i, \"\"))\n", + " print(s)\n", " print(f\"===== + Honesty Control =====\")\n", - " print(p['generated_text'][0].replace(i, \"\"))\n", + " print(p)\n", " print()\n", " print(f\"===== - Honesty Control =====\")\n", - " print(n['generated_text'][0].replace(i, \"\"))\n", + " print(n)\n", " print()\n" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 71, "metadata": {}, - "outputs": [], - "source": [] + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "===== No Control =====\n", + "2.03%\n", + "===== + Honesty Control =====\n", + "23.10%\n", + "\n", + "===== - Honesty Control =====\n", + "1.02%\n", + "\n", + "===== No Control =====\n", + "33.63%\n", + "===== + Honesty Control =====\n", + "83.52%\n", + "\n", + "===== - Honesty Control =====\n", + "7.75%\n", + "\n" + ] + } + ], + "source": [ + "for i,s,p,n in zip(inputs, baseline_outputs['ans'], control_outputs['ans'], control_outputs_neg['ans']):\n", + " print(\"===== No Control =====\")\n", + " print(f'{s:02.2%}')\n", + " print(f\"===== + Honesty Control =====\")\n", + " print(f'{p:02.2%}')\n", + " print()\n", + " print(f\"===== - Honesty Control =====\")\n", + " print(f'{n:02.2%}')\n", + " print()\n" + ] }, { "cell_type": "code", diff --git a/src/repe/pipelines.py b/src/repe/pipelines.py index ababfbf..0c08c56 100644 --- a/src/repe/pipelines.py +++ b/src/repe/pipelines.py @@ -1,7 +1,8 @@ from transformers import AutoModel, AutoModelForCausalLM from transformers.pipelines import PIPELINE_REGISTRY from .rep_reading_pipeline import RepReadingPipeline -from .rep_control_pipeline import RepControlPipeline, RepControlPipeline2 +from .rep_control_pipeline import RepControlPipeline +from .rep_control_pipeline_baukit import RepControlPipeline2 def repe_pipeline_registry(): PIPELINE_REGISTRY.register_pipeline( diff --git a/src/repe/rep_control_pipeline.py b/src/repe/rep_control_pipeline.py index e07f450..c89645e 100644 --- a/src/repe/rep_control_pipeline.py +++ b/src/repe/rep_control_pipeline.py @@ -7,9 +7,11 @@ from transformers.pipelines import ( from transformers.pipelines.base import GenericTensor from .rep_control_reading_vec import WrappedReadingVecModel from typing import Dict +from .rep_control_pipeline_baukit import RepControlPipeline2 -class RepControlPipeline(FeatureExtractionPipeline): + +class RepControlPipeline(RepControlPipeline2): def __init__( self, model, @@ -35,12 +37,6 @@ class RepControlPipeline(FeatureExtractionPipeline): 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() @@ -50,112 +46,3 @@ class RepControlPipeline(FeatureExtractionPipeline): 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_pipeline_baukit.py b/src/repe/rep_control_pipeline_baukit.py new file mode 100644 index 0000000..3593d33 --- /dev/null +++ b/src/repe/rep_control_pipeline_baukit.py @@ -0,0 +1,109 @@ +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 List, Tuple, Dict, Any, Union, NewType +from baukit.nethook import Trace, TraceDict, recursive_copy +from functools import partial +from src.datasets.scores import choice2ids +from src.datasets.scores import scores2choice_probs + + +Activations = NewType("InterventionDict", Dict[str, torch.Tensor]) + +def row_choice_ids(answer_choices, tokenizer): + return choice2ids([[c] for c in answer_choices], tokenizer) + + +def intervene(output, activation): + assert output.ndim == 3, f"expected output to be (batch, seq, vocab), got {output.shape}" + return output + activation.to(output.device)[None, None, :] + +def intervention_meta_fn2( + outputs: 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: + ... + + """ + if type(outputs) is tuple: + output0 = intervene(outputs[0], activations[layer_name]) + return tuple(output0, *outputs[1:]) + elif type(outputs) is torch.Tensor: + return intervene(outputs, activations[layer_name]) + else: + raise ValueError(f"outputs must be tuple or tensor, got {type(outputs)}") + + +class RepControlPipeline2(FeatureExtractionPipeline): + """This version uses baukit.""" + def __init__(self, model, tokenizer, max_length, layer_name_tmpl="model.layers.{}", **kwargs): + super().__init__(model=model, tokenizer=tokenizer, **kwargs) + self.max_length = max_length + self.layer_name_tmpl = layer_name_tmpl + + def __call__(self, model_inputs, activations=None, **kwargs): + if activations is not None: + activations_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations.items()}) + layers_names = [self.layer_name_tmpl.format(i) for i in activations.keys()] + edit_fn = partial(intervention_meta_fn2, activations=activations_i) + 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(ac, self.tokenizer) for ac in o['answer_choices']] + + p = o['add_ans'] = torch.stack([scores2choice_probs2(l, c) for l, c in zip(o['end_logits'], o['choice_ids'])]) + o['ans'] = p[:, 1] / (torch.sum(p, 1) + 1e-5) + + return o + +def scores2choice_probs2(logits, choiceids: List[List[int]]): + """calculate the probability for each group of choices.""" + assert logits.ndim==1, f"expected logits to be 1d, got {logits.shape}" + probs = torch.softmax(logits, 0) # shape [tokens, inferences) + probs_c = torch.tensor([[probs[cc] for cc in c] for c in choiceids]).sum(1) # sum over alternate choices e.g. [['decrease', 'dec'],['inc', 'increase']] + return probs_c