diff --git a/mjc_notes.md b/mjc_notes.md index ea341e8..4442a26 100644 --- a/mjc_notes.md +++ b/mjc_notes.md @@ -2169,7 +2169,7 @@ best to use a class.... and pass it around... doesn't need to be a pipeline but none of them seem to have a reasonable magnitude so.... not sure if any will give valid ones -```sh +~~~sh export ORIGINAL_ORG=TheBloke export NEW_ORG=wassname export MODEL_NAME=phi-2-GPTQ @@ -2184,7 +2184,9 @@ git remote add upstream https://huggingface.co/$ORIGINAL_ORG/$MODEL_NAME git fetch upstream git rebase upstream/main git push --force-with-lease -``` +~~~ + + ok I wasn't even applying the activation right? I was adding, mean to multiply... Definitly time to simplify @@ -2195,7 +2197,7 @@ plan - src.datasets.intervene create_cache_interventions -# Phi-2 +## Phi-2 scratch 2023-12-16 12:43:50 model = AutoModelForCausalLM.from_pretrained(ckpt_path, torch_dtype=torch.float16, flash_attn=True, flash_rotary=True, fused_dense=True) @@ -2207,3 +2209,36 @@ maybe for padding use 50256? Rather than 0? "transformers_version": "4.37.0.dev0", +```py +com_directions = get_com_directions(num_layers, num_heads, train_set_idxs, val_set_idxs, separated_head_wise_activations, separated_labels) +direction = com_directions[layer_head_to_flattened_idx(layer, head, num_heads)] +direction = direction / np.linalg.norm(direction) + +interventions = {} +activations = tuning_activations[:,layer,head,:] # batch x 128 +proj_vals = activations @ direction.T +proj_val_std = np.std(proj_vals) +interventions[f"model.layers.{layer}.self_attn.head_out"].append((head, direction.squeeze(), proj_val_std)) + +head_output[:, -1, head, :] += args.alpha * proj_val_std * direction_to_add +``` + + +Oh wait the actual intervention in geometry of truth is + +``` +direction = direction / direction.norm() +diff = (true_mean - false_mean) @ direction +direction = diff * direction + +# wtf +output[0][:, intervention_idx, :] += direction * alpha +``` + +# 2023-12-16 16:50:20 + +OK I would like a sklearn like interface + +```py +class Intervention(nn.Module): + def _ diff --git a/notebooks/104_check_intervention.ipynb b/notebooks/104_check_intervention.ipynb index fd799c5..d394e3c 100644 --- a/notebooks/104_check_intervention.ipynb +++ b/notebooks/104_check_intervention.ipynb @@ -9,28 +9,9 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "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" - ] - }, - { - "data": { - "text/plain": [ - "1" - ] - }, - "execution_count": 1, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "\n", "import os\n", @@ -58,18 +39,9 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "CUDA extension not installed.\n", - "CUDA extension not installed.\n" - ] - } - ], + "outputs": [], "source": [ "# load my code\n", "%load_ext autoreload\n", @@ -94,7 +66,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -117,98 +89,9 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "ExtractConfig(datasets=('amazon_polarity', 'super_glue:boolq', 'glue:qnli', 'imdb'), model='wassname/phi-2-GPTQ_w_hidden_states', batch_size=5, pad_token_id=50256, prompt_format='phi', data_dirs=(), max_examples=(400, 400), num_shots=2, num_variants=-1, seed=42, template_path=None, max_length=1000, disable_ds_cache=False, intervention_direction_method='cluster_mean', intervention_fit_examples=160, intervention_layer_name_template='transformer.h.{}')\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "\u001b[32m2023-12-16 11:51:16.540\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m24\u001b[0m - \u001b[1mchanging use_cache from True to False\u001b[0m\n", - "2023-12-16T11:51:16.540903+0800 INFO changing use_cache from True to False\n", - "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n", - "\u001b[32m2023-12-16 11:51:16.902\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m24\u001b[0m - \u001b[1mchanging pad_token_id from None to 50256\u001b[0m\n", - "2023-12-16T11:51:16.902047+0800 INFO changing pad_token_id from None to 50256\n", - "\u001b[32m2023-12-16 11:51:16.902\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m24\u001b[0m - \u001b[1mchanging padding_side from right to left\u001b[0m\n", - "2023-12-16T11:51:16.902868+0800 INFO changing padding_side from right to left\n", - "\u001b[32m2023-12-16 11:51:16.903\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m24\u001b[0m - \u001b[1mchanging truncation_side from right to left\u001b[0m\n", - "2023-12-16T11:51:16.903473+0800 INFO changing truncation_side from right to left\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "PhiForCausalLM(\n", - " (transformer): PhiModel(\n", - " (embd): Embedding(\n", - " (wte): Embedding(51200, 2560)\n", - " (drop): Dropout(p=0.0, inplace=False)\n", - " )\n", - " (h): ModuleList(\n", - " (0-31): 32 x ParallelBlock(\n", - " (ln): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n", - " (resid_dropout): Dropout(p=0.1, inplace=False)\n", - " (mixer): MHA(\n", - " (rotary_emb): RotaryEmbedding()\n", - " (inner_attn): SelfAttention(\n", - " (drop): Dropout(p=0.0, inplace=False)\n", - " )\n", - " (inner_cross_attn): CrossAttention(\n", - " (drop): Dropout(p=0.0, inplace=False)\n", - " )\n", - " (Wqkv): QuantLinear()\n", - " (out_proj): QuantLinear()\n", - " )\n", - " (mlp): MLP(\n", - " (act): NewGELUActivation()\n", - " (fc1): QuantLinear()\n", - " (fc2): QuantLinear()\n", - " )\n", - " )\n", - " )\n", - " )\n", - " (lm_head): CausalLMHead(\n", - " (ln): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n", - " (linear): Linear(in_features=2560, out_features=51200, bias=True)\n", - " )\n", - " (loss): CausalLMLoss(\n", - " (loss_fct): CrossEntropyLoss()\n", - " )\n", - ")\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Generating train split: 0 examples [00:00, ? examples/s]\u001b[32m2023-12-16 11:51:22.526\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.prompts.prompt_loading\u001b[0m:\u001b[36mload_prompts\u001b[0m:\u001b[36m120\u001b[0m - \u001b[1mExtracting 11 variants of each prompt\u001b[0m\n", - "2023-12-16T11:51:22.526803+0800 INFO Extracting 11 variants of each prompt\n", - "Generating train split: 2402 examples [05:52, 6.82 examples/s]\n", - "format_prompt: 100%|██████████| 2402/2402 [00:00<00:00, 7939.78 examples/s]\n", - "tokenize: 100%|██████████| 2402/2402 [00:01<00:00, 1312.50 examples/s]\n", - "truncated: 100%|██████████| 2402/2402 [00:00<00:00, 2617.70 examples/s]\n", - "truncated: 100%|██████████| 2402/2402 [00:00<00:00, 2537.01 examples/s]\n", - "prompt_truncated: 100%|██████████| 2402/2402 [00:07<00:00, 322.85 examples/s]\n", - "choice_ids: 100%|██████████| 2402/2402 [00:00<00:00, 4870.83 examples/s]\n", - "\u001b[32m2023-12-16 11:57:23.152\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.prompts.prompt_loading\u001b[0m:\u001b[36mload_preproc_dataset\u001b[0m:\u001b[36m364\u001b[0m - \u001b[1mmedian token length: 433.0 for amazon_polarity. max_length=1000\u001b[0m\n", - "2023-12-16T11:57:23.152948+0800 INFO median token length: 433.0 for amazon_polarity. max_length=1000\n", - "\u001b[32m2023-12-16 11:57:23.154\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.prompts.prompt_loading\u001b[0m:\u001b[36mload_preproc_dataset\u001b[0m:\u001b[36m368\u001b[0m - \u001b[1mtruncation rate: 0.00% on amazon_polarity\u001b[0m\n", - "2023-12-16T11:57:23.154377+0800 INFO truncation rate: 0.00% on amazon_polarity\n", - "Filter: 100%|██████████| 2402/2402 [00:01<00:00, 2283.72 examples/s]\n", - "Filter: 100%|██████████| 2402/2402 [00:01<00:00, 2201.41 examples/s]\n", - "\u001b[32m2023-12-16 11:57:25.313\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.prompts.prompt_loading\u001b[0m:\u001b[36mload_preproc_dataset\u001b[0m:\u001b[36m377\u001b[0m - \u001b[1mnum_rows (after filtering out truncated rows) 2402=>2402\u001b[0m\n", - "2023-12-16T11:57:25.313993+0800 INFO num_rows (after filtering out truncated rows) 2402=>2402\n" - ] - } - ], + "outputs": [], "source": [ "ds_name='amazon_polarity'\n", "cfg = ExtractConfig(max_examples=(400, 400),\n", @@ -241,28 +124,9 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "\u001b[32m2023-12-16 11:57:25.647\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.datasets.intervene\u001b[0m:\u001b[36mcreate_cache_interventions\u001b[0m:\u001b[36m138\u001b[0m - \u001b[1mLoaded interventions from /media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/data/interventions/wassname-phi-2-GPTQ_w_hidden_states_+_cluster_mean_481.pkl\u001b[0m\n", - "2023-12-16T11:57:25.647973+0800 INFO Loaded interventions from /media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/data/interventions/wassname-phi-2-GPTQ_w_hidden_states_+_cluster_mean_481.pkl\n" - ] - }, - { - "data": { - "text/plain": [ - "" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "honesty_rep_reader = create_cache_interventions(model, tokenizer, cfg)\n", "honesty_rep_reader\n" @@ -270,41 +134,18 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "" - ] - }, - "execution_count": 6, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "honesty_rep_reader\n" ] }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "array([[-0.1623942 , -0.05915324, 0.08666992, ..., 0.04361978,\n", - " 0.13239542, 0.11653644]], dtype=float32)" - ] - }, - "execution_count": 7, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "layer = 6\n", "honesty_rep_reader.directions[layer]\n" @@ -312,45 +153,9 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{6: [-1],\n", - " 7: [-1],\n", - " 8: [-1],\n", - " 9: [-1],\n", - " 10: [-1],\n", - " 11: [-1],\n", - " 12: [-1],\n", - " 13: [-1],\n", - " 14: [-1],\n", - " 15: [-1],\n", - " 16: [-1],\n", - " 17: [-1],\n", - " 18: [-1],\n", - " 19: [-1],\n", - " 20: [-1],\n", - " 21: [-1],\n", - " 22: [-1],\n", - " 23: [-1],\n", - " 24: [-1],\n", - " 25: [-1],\n", - " 26: [-1],\n", - " 27: [-1],\n", - " 28: [-1],\n", - " 29: [-1],\n", - " 30: [-1],\n", - " 31: [-1]}" - ] - }, - "execution_count": 8, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "honesty_rep_reader.direction_signs\n" ] @@ -371,30 +176,9 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "rep-reading is already registered. Overwriting pipeline for task rep-reading...\n", - "rep-control2 is already registered. Overwriting pipeline for task rep-control2...\n", - "\u001b[32m2023-12-16 11:57:25.760\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.datasets.intervene\u001b[0m:\u001b[36mcreate_cache_interventions\u001b[0m:\u001b[36m138\u001b[0m - \u001b[1mLoaded interventions from /media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/data/interventions/wassname-phi-2-GPTQ_w_hidden_states_+_cluster_mean_481.pkl\u001b[0m\n", - "2023-12-16T11:57:25.760677+0800 INFO Loaded interventions from /media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/data/interventions/wassname-phi-2-GPTQ_w_hidden_states_+_cluster_mean_481.pkl\n" - ] - }, - { - "data": { - "text/plain": [ - "" - ] - }, - "execution_count": 9, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "\n", "from src.repe import repe_pipeline_registry\n", @@ -418,14 +202,13 @@ " \"rep-control2\", \n", " model=model, \n", " tokenizer=tokenizer, \n", - " layers=hidden_layers, \n", - " max_length=cfg.max_length, layer_name_tmpl=cfg.intervention_layer_name_template)\n", + " layers=hidden_layers)\n", "rep_control_pipeline2\n" ] }, { "cell_type": "code", - "execution_count": 23, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -472,39 +255,9 @@ }, { "cell_type": "code", - "execution_count": 25, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "ename": "RuntimeError", - "evalue": "1D tensors expected, but got 3D and 1D tensors", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)", - "\u001b[1;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/notebooks/104_check_intervention.ipynb Cell 18\u001b[0m line \u001b[0;36m5\n\u001b[1;32m 1\u001b[0m ds \u001b[39m=\u001b[39m ds_tokens\u001b[39m.\u001b[39mselect(\u001b[39mrange\u001b[39m(\u001b[39m3\u001b[39m))\u001b[39m.\u001b[39mto_iterable_dataset()\n\u001b[1;32m 2\u001b[0m r1 \u001b[39m=\u001b[39m rep_control_pipeline2(model_inputs\u001b[39m=\u001b[39mds,\n\u001b[1;32m 3\u001b[0m activations\u001b[39m=\u001b[39mactivations_neg_i,\n\u001b[1;32m 4\u001b[0m batch_size\u001b[39m=\u001b[39mbatch_size,)\n\u001b[0;32m----> 5\u001b[0m r \u001b[39m=\u001b[39m \u001b[39mlist\u001b[39;49m(r1)\n\u001b[1;32m 6\u001b[0m o \u001b[39m=\u001b[39m r[\u001b[39m0\u001b[39m]\n\u001b[1;32m 7\u001b[0m print_pipeline_r(o)\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/transformers/pipelines/pt_utils.py:124\u001b[0m, in \u001b[0;36mPipelineIterator.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 121\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mloader_batch_item()\n\u001b[1;32m 123\u001b[0m \u001b[39m# We're out of items within a batch\u001b[39;00m\n\u001b[0;32m--> 124\u001b[0m item \u001b[39m=\u001b[39m \u001b[39mnext\u001b[39;49m(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49miterator)\n\u001b[1;32m 125\u001b[0m processed \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39minfer(item, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mparams)\n\u001b[1;32m 126\u001b[0m \u001b[39m# We now have a batch of \"inferred things\".\u001b[39;00m\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/transformers/pipelines/pt_utils.py:125\u001b[0m, in \u001b[0;36mPipelineIterator.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 123\u001b[0m \u001b[39m# We're out of items within a batch\u001b[39;00m\n\u001b[1;32m 124\u001b[0m item \u001b[39m=\u001b[39m \u001b[39mnext\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39miterator)\n\u001b[0;32m--> 125\u001b[0m processed \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49minfer(item, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mparams)\n\u001b[1;32m 126\u001b[0m \u001b[39m# We now have a batch of \"inferred things\".\u001b[39;00m\n\u001b[1;32m 127\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mloader_batch_size \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 128\u001b[0m \u001b[39m# Try to infer the size of the batch\u001b[39;00m\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/transformers/pipelines/base.py:1046\u001b[0m, in \u001b[0;36mPipeline.forward\u001b[0;34m(self, model_inputs, **forward_params)\u001b[0m\n\u001b[1;32m 1044\u001b[0m \u001b[39mwith\u001b[39;00m inference_context():\n\u001b[1;32m 1045\u001b[0m model_inputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_ensure_tensor_on_device(model_inputs, device\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdevice)\n\u001b[0;32m-> 1046\u001b[0m model_outputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_forward(model_inputs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mforward_params)\n\u001b[1;32m 1047\u001b[0m model_outputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_ensure_tensor_on_device(model_outputs, device\u001b[39m=\u001b[39mtorch\u001b[39m.\u001b[39mdevice(\u001b[39m\"\u001b[39m\u001b[39mcpu\u001b[39m\u001b[39m\"\u001b[39m))\n\u001b[1;32m 1048\u001b[0m \u001b[39melse\u001b[39;00m:\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/src/repe/rep_control_pipeline_baukit.py:124\u001b[0m, in \u001b[0;36mRepControlPipeline2._forward\u001b[0;34m(self, inputs, activations)\u001b[0m\n\u001b[1;32m 120\u001b[0m \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39mno_grad():\n\u001b[1;32m 121\u001b[0m \u001b[39mwith\u001b[39;00m TraceDict(\n\u001b[1;32m 122\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel, layers_names, detach\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m, edit_output\u001b[39m=\u001b[39medit_fn_pos\n\u001b[1;32m 123\u001b[0m ) \u001b[39mas\u001b[39;00m ret:\n\u001b[0;32m--> 124\u001b[0m outputs_pos \u001b[39m=\u001b[39m transform_model_output(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mmodel(\u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mmodel_in))\n\u001b[1;32m 126\u001b[0m \u001b[39mwith\u001b[39;00m TraceDict(\n\u001b[1;32m 127\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel, layers_names, detach\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m, edit_output\u001b[39m=\u001b[39medit_fn_neu\n\u001b[1;32m 128\u001b[0m ) \u001b[39mas\u001b[39;00m ret:\n\u001b[1;32m 129\u001b[0m outputs_neg \u001b[39m=\u001b[39m transform_model_output(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel(\u001b[39m*\u001b[39m\u001b[39m*\u001b[39mmodel_in))\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1525\u001b[0m \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1526\u001b[0m \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 1530\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/accelerate/hooks.py:165\u001b[0m, in \u001b[0;36madd_hook_to_module..new_forward\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 163\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[1;32m 164\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 165\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 166\u001b[0m \u001b[39mreturn\u001b[39;00m module\u001b[39m.\u001b[39m_hf_hook\u001b[39m.\u001b[39mpost_forward(module, output)\n", - "File \u001b[0;32m~/.cache/huggingface/modules/transformers_modules/wassname/phi-2-GPTQ_w_hidden_states/e39f09ffdefdb011e32b1cad1175c952a5a0f88d/modeling_phi.py:957\u001b[0m, in \u001b[0;36mPhiForCausalLM.forward\u001b[0;34m(self, input_ids, past_key_values, attention_mask, labels, **kwargs)\u001b[0m\n\u001b[1;32m 949\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\n\u001b[1;32m 950\u001b[0m \u001b[39mself\u001b[39m,\n\u001b[1;32m 951\u001b[0m input_ids: torch\u001b[39m.\u001b[39mLongTensor,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 955\u001b[0m \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs,\n\u001b[1;32m 956\u001b[0m ) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m CausalLMOutputWithPast:\n\u001b[0;32m--> 957\u001b[0m outputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mtransformer(input_ids, past_key_values\u001b[39m=\u001b[39;49mpast_key_values, attention_mask\u001b[39m=\u001b[39;49mattention_mask)\n\u001b[1;32m 958\u001b[0m lm_logits \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlm_head(outputs\u001b[39m.\u001b[39mlast_hidden_state)\n\u001b[1;32m 960\u001b[0m loss \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1525\u001b[0m \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1526\u001b[0m \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 1530\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n", - "File \u001b[0;32m~/.cache/huggingface/modules/transformers_modules/wassname/phi-2-GPTQ_w_hidden_states/e39f09ffdefdb011e32b1cad1175c952a5a0f88d/modeling_phi.py:917\u001b[0m, in \u001b[0;36mPhiModel.forward\u001b[0;34m(self, input_ids, past_key_values, attention_mask)\u001b[0m\n\u001b[1;32m 915\u001b[0m all_hidden_states \u001b[39m=\u001b[39m [hidden_states]\n\u001b[1;32m 916\u001b[0m \u001b[39mfor\u001b[39;00m layer \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mh:\n\u001b[0;32m--> 917\u001b[0m hidden_states, attn_outputs \u001b[39m=\u001b[39m layer_outputs \u001b[39m=\u001b[39m layer(\n\u001b[1;32m 918\u001b[0m hidden_states,\n\u001b[1;32m 919\u001b[0m past_key_values\u001b[39m=\u001b[39;49mpast_key_values,\n\u001b[1;32m 920\u001b[0m attention_mask\u001b[39m=\u001b[39;49mattention_mask,\n\u001b[1;32m 921\u001b[0m )\n\u001b[1;32m 922\u001b[0m all_hidden_states\u001b[39m.\u001b[39mappend(hidden_states)\n\u001b[1;32m 923\u001b[0m all_self_attns\u001b[39m.\u001b[39mappend(attn_outputs)\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1581\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1579\u001b[0m hook_result \u001b[39m=\u001b[39m hook(\u001b[39mself\u001b[39m, args, kwargs, result)\n\u001b[1;32m 1580\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1581\u001b[0m hook_result \u001b[39m=\u001b[39m hook(\u001b[39mself\u001b[39;49m, args, result)\n\u001b[1;32m 1583\u001b[0m \u001b[39mif\u001b[39;00m hook_result \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 1584\u001b[0m result \u001b[39m=\u001b[39m hook_result\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/baukit/nethook.py:73\u001b[0m, in \u001b[0;36mTrace.__init__..retain_hook\u001b[0;34m(m, inputs, output)\u001b[0m\n\u001b[1;32m 71\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mretain_hook\u001b[39m(m, inputs, output):\n\u001b[1;32m 72\u001b[0m \u001b[39mif\u001b[39;00m edit_output:\n\u001b[0;32m---> 73\u001b[0m output \u001b[39m=\u001b[39m invoke_with_optional_args(\n\u001b[1;32m 74\u001b[0m edit_output, output\u001b[39m=\u001b[39;49moutput, layer\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mlayer, inputs\u001b[39m=\u001b[39;49minputs\n\u001b[1;32m 75\u001b[0m )\n\u001b[1;32m 76\u001b[0m \u001b[39mif\u001b[39;00m retain_input:\n\u001b[1;32m 77\u001b[0m retainer\u001b[39m.\u001b[39minput \u001b[39m=\u001b[39m recursive_copy(\n\u001b[1;32m 78\u001b[0m inputs[\u001b[39m0\u001b[39m] \u001b[39mif\u001b[39;00m \u001b[39mlen\u001b[39m(inputs) \u001b[39m==\u001b[39m \u001b[39m1\u001b[39m \u001b[39melse\u001b[39;00m inputs,\n\u001b[1;32m 79\u001b[0m clone\u001b[39m=\u001b[39mclone,\n\u001b[1;32m 80\u001b[0m detach\u001b[39m=\u001b[39mdetach,\n\u001b[1;32m 81\u001b[0m retain_grad\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m,\n\u001b[1;32m 82\u001b[0m ) \u001b[39m# retain_grad applies to output only.\u001b[39;00m\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/baukit/nethook.py:471\u001b[0m, in \u001b[0;36minvoke_with_optional_args\u001b[0;34m(fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m 469\u001b[0m \u001b[39mif\u001b[39;00m argspec\u001b[39m.\u001b[39mvarargs \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 470\u001b[0m pass_args \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39mlist\u001b[39m(args[used_pos:])\n\u001b[0;32m--> 471\u001b[0m \u001b[39mreturn\u001b[39;00m fn(\u001b[39m*\u001b[39;49mpass_args, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mpass_kw)\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/src/datasets/intervene.py:52\u001b[0m, in \u001b[0;36mintervention_meta_fn2\u001b[0;34m(outputs, layer_name, activations)\u001b[0m\n\u001b[1;32m 41\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"see\u001b[39;00m\n\u001b[1;32m 42\u001b[0m \u001b[39m- honest_llama: https://github.com/likenneth/honest_llama/blob/e010f82bfbeaa4326cef8493b0dd5b8b14c6da67/validation/validate_2fold.py#L114\u001b[39;00m\n\u001b[1;32m 43\u001b[0m \u001b[39m- baukit: https://github.com/davidbau/baukit/blob/main/baukit/nethook.py#L42C1-L45C56\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 48\u001b[0m \u001b[39m ...\u001b[39;00m\n\u001b[1;32m 49\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 50\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mtype\u001b[39m(outputs) \u001b[39mis\u001b[39;00m \u001b[39mtuple\u001b[39m:\n\u001b[1;32m 51\u001b[0m \u001b[39m# just edit the first one, and put it back in the tuple\u001b[39;00m\n\u001b[0;32m---> 52\u001b[0m output0 \u001b[39m=\u001b[39m intervene(outputs[\u001b[39m0\u001b[39;49m], activations[layer_name])\n\u001b[1;32m 53\u001b[0m \u001b[39mreturn\u001b[39;00m (output0, \u001b[39m*\u001b[39moutputs[\u001b[39m1\u001b[39m:])\n\u001b[1;32m 54\u001b[0m \u001b[39melif\u001b[39;00m \u001b[39mtype\u001b[39m(outputs) \u001b[39mis\u001b[39;00m torch\u001b[39m.\u001b[39mTensor:\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/src/datasets/intervene.py:33\u001b[0m, in \u001b[0;36mintervene\u001b[0;34m(output, activation)\u001b[0m\n\u001b[1;32m 30\u001b[0m \u001b[39m# assert torch.isfinite(output).all(), 'model output nan'\u001b[39;00m\n\u001b[1;32m 31\u001b[0m output2 \u001b[39m=\u001b[39m output \u001b[39m+\u001b[39m activation\u001b[39m.\u001b[39mto(output\u001b[39m.\u001b[39mdevice)[\u001b[39mNone\u001b[39;00m, :]\n\u001b[0;32m---> 33\u001b[0m output2 \u001b[39m=\u001b[39m project_onto_direction(output, activation)\n\u001b[1;32m 34\u001b[0m \u001b[39m# assert torch.isfinite(output2).all(), 'intervention lead to nan'\u001b[39;00m\n\u001b[1;32m 35\u001b[0m \u001b[39mreturn\u001b[39;00m output2\n", - "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/src/repe/rep_readers.py:15\u001b[0m, in \u001b[0;36mproject_onto_direction\u001b[0;34m(H, direction)\u001b[0m\n\u001b[1;32m 13\u001b[0m mag \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mlinalg\u001b[39m.\u001b[39mnorm(direction)\u001b[39m.\u001b[39mto(H\u001b[39m.\u001b[39mdevice)\n\u001b[1;32m 14\u001b[0m \u001b[39massert\u001b[39;00m torch\u001b[39m.\u001b[39misfinite(mag)\n\u001b[0;32m---> 15\u001b[0m \u001b[39mreturn\u001b[39;00m H\u001b[39m.\u001b[39;49mdot(direction) \u001b[39m/\u001b[39m mag\n", - "\u001b[0;31mRuntimeError\u001b[0m: 1D tensors expected, but got 3D and 1D tensors" - ] - } - ], + "outputs": [], "source": [ "ds = ds_tokens.select(range(3)).to_iterable_dataset()\n", "r1 = rep_control_pipeline2(model_inputs=ds,\n", diff --git a/notebooks/make_dataset2.py b/notebooks/make_dataset2.py index 0475c1b..6c10555 100644 --- a/notebooks/make_dataset2.py +++ b/notebooks/make_dataset2.py @@ -103,19 +103,13 @@ if TEST: # Fit an intervention -honesty_rep_reader1 = create_cache_interventions( +intervention = create_cache_interventions( model, tokenizer, cfg, ) -# honesty_rep_reader2 = create_cache_interventions( -# model, -# tokenizer, -# cfg, -# ) - -hidden_layers = sorted(honesty_rep_reader1.directions.keys()) +hidden_layers = sorted(intervention.direction.keys()) hidden_layers @@ -140,21 +134,15 @@ rep_control_pipeline2 = pipeline( tokenizer=tokenizer, layers=hidden_layers, max_length=cfg.max_length, - layer_name_tmpl=cfg.intervention_layer_name_template ) rep_control_pipeline2 # %% -from src.datasets.intervene import get_activations_from_reader +# from src.datasets.intervene import get_activations_from_reader from src.datasets.intervene import test_intervention_quality -activations = get_activations_from_reader( - honesty_rep_reader1, hidden_layers, dtype=model.dtype, device=model.device -) -# activations2 = get_activations_from_reader( -# honesty_rep_reader2, hidden_layers, dtype=model.dtype, device=model.device -# ) + # %% @@ -168,7 +156,7 @@ def create_hs_ds( ds_name, ds_tokens, pipeline, - activations=None, + intervention=None, f=None, batch_size=2, split_type="train", @@ -212,7 +200,7 @@ def create_hs_ds( # first we make the calibration dataset with no intervention gen_kwargs = dict( model_inputs=ds, - activations=activations, + intervention=intervention, batch_size=batch_size, ) @@ -273,7 +261,7 @@ for ds_name in cfg.datasets: assert len(dataset_test) > 3 # FIXME: - test_intervention_quality(dataset_train, activations, model, rep_control_pipeline2, batch_size=batch_size, ds_name=ds_name) + test_intervention_quality(dataset_train, intervention, model, rep_control_pipeline2, batch_size=batch_size, ds_name=ds_name) ds1, f = create_hs_ds( ds_name, @@ -282,7 +270,7 @@ for ds_name in cfg.datasets: split_type="train", debug=True, batch_size=batch_size, - activations=activations, + intervention=intervention, ) clear_mem() ds1, f = create_hs_ds( @@ -292,7 +280,7 @@ for ds_name in cfg.datasets: split_type="test", debug=True, batch_size=batch_size, - activations=activations, + intervention=intervention, ) clear_mem() diff --git a/poetry.lock b/poetry.lock index 059fe98..6cc8936 100644 --- a/poetry.lock +++ b/poetry.lock @@ -5025,4 +5025,4 @@ cffi = ["cffi (>=1.11)"] [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.13" -content-hash = "01142df3e809581847ae1c009141e078a9a7440a87e768559a31c4aafae6e42d" +content-hash = "c8053335e646105c66b57dac71353d058145f69e260705844745961df43fa210" diff --git a/src/datasets/intervene.py b/src/datasets/intervene.py index 1314d9b..e6defda 100644 --- a/src/datasets/intervene.py +++ b/src/datasets/intervene.py @@ -13,48 +13,49 @@ from src.config import root_folder from src.prompts.prompt_loading import load_preproc_dataset from transformers import AutoTokenizer, pipeline, Pipeline from loguru import logger -from src.repe.rep_readers import project_onto_direction +from jaxtyping import Float +from torch import nn, Tensor -Activations = NewType("Activations", Dict[str, torch.Tensor]) +# Activations = NewType("Activations", Dict[str, torch.Tensor]) -InterventionDict = NewType( - "InterventionDict", Dict[str, List[Tuple[np.ndarray, float]]] -) +# InterventionDict = NewType( +# "InterventionDict", Dict[str, List[Tuple[np.ndarray, float]]] +# ) -def intervene(output, activation): - # TODO need attention mask - assert ( - output.ndim == 3 - ), f"expected output to be (batch, seq, vocab), got {output.shape}" - # assert torch.isfinite(output).all(), 'model output nan' - output2 = output + activation.to(output.device)[None, :] +# def intervene(output, activation): +# # TODO need attention mask +# assert ( +# output.ndim == 3 +# ), f"expected output to be (batch, seq, vocab), got {output.shape}" +# # assert torch.isfinite(output).all(), 'model output nan' +# output2 = output + activation.to(output.device)[None, :] - output2 = project_onto_direction(output, activation) - # assert torch.isfinite(output2).all(), 'intervention lead to nan' - return output2 +# output2 = project_onto_direction(output, activation) +# # assert torch.isfinite(output2).all(), 'intervention lead to nan' +# return output2 -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 +# 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: - # just edit the first one, and put it back in the tuple - output0 = intervene(outputs[0], activations[layer_name]) - return (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)}") +# 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: +# # just edit the first one, and put it back in the tuple +# output0 = intervene(outputs[0], activations[layer_name]) +# return (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)}") def create_cache_interventions( @@ -113,7 +114,7 @@ def create_cache_interventions( train_labels = train_labels == 0 rep_reading_pipeline = pipeline("rep-reading", model=model, tokenizer=tokenizer) - honesty_rep_reader = rep_reading_pipeline.get_directions( + intervention = rep_reading_pipeline.get_directions( dataset_fit["question"], rep_token=rep_token, hidden_layers=hidden_layers, @@ -121,28 +122,29 @@ def create_cache_interventions( train_labels=dataset_fit["label_true"], direction_method=direction_method, batch_size=batch_size, + layer_name_tmpl=cfg.intervention_layer_name_template, **tokenizer_args, ) assert np.isfinite( - np.concatenate(list(honesty_rep_reader.directions.values())) + np.concatenate(list(intervention.direction.values())) ).all() # assert torch.isfinite(torch.concat(list(honesty_rep_reader.directions.values()))).all() # and save with open(intervention_f, "wb") as f: - pickle.dump(honesty_rep_reader, f) + pickle.dump(intervention, f) logger.info(f"Saved interventions to {intervention_f}") with open(intervention_f, "rb") as f: - honesty_rep_reader = pickle.load(f) + intervention = pickle.load(f) logger.info(f"Loaded interventions from {intervention_f}") - return honesty_rep_reader + return intervention def test_intervention_quality( - dataset_train, activations, model, rep_control_pipeline2, batch_size=2, ds_name="" + dataset_train, intervention, model, rep_control_pipeline2, batch_size=2, ds_name="" ): """ Check interventions are ordered and different and valid @@ -155,7 +157,7 @@ def test_intervention_quality( batch = inputs[batch_index * batch_size : (batch_index + 1) * batch_size] with torch.no_grad(): baseline_outputs += rep_control_pipeline2( - batch, batch_size=batch_size, activations=activations + batch, batch_size=batch_size, intervention=intervention ) # So here we check that the interventions are ordered, e.g. the positive one gives a more true answer than the neutral or negative ones @@ -205,24 +207,22 @@ def test_intervention_quality( return df -def get_activations_from_reader( - honesty_rep_reader: Pipeline, hidden_layers: list, coeff=1, dtype=None, device=None -) -> Dict[str, float]: - """Get activations from the honesty_rep_reader""" +# def get_activations_from_reader( +# honesty_rep_reader: Pipeline, hidden_layers: list, coeff=1, dtype=None, device=None +# ) -> Dict[str, float]: +# """Get activations from the honesty_rep_reader""" - # FIXME: coeff is a magic number. The representation_engineering repo used 8, but it seems to vary by model? +# activations = {} +# for layer in hidden_layers: +# activations[layer] = torch.tensor( +# coeff +# * honesty_rep_reader.directions[layer] +# * honesty_rep_reader.direction_signs[layer] +# ) +# if device: +# activations[layer] = activations[layer].to(device) +# if dtype: +# activations[layer] = activations[layer].to(dtype) - activations = {} - for layer in hidden_layers: - activations[layer] = torch.tensor( - coeff - * honesty_rep_reader.directions[layer] - * honesty_rep_reader.direction_signs[layer] - ) - if device: - activations[layer] = activations[layer].to(device) - if dtype: - activations[layer] = activations[layer].to(dtype) - - assert torch.isfinite(torch.concat(list(activations.values()))).all() - return activations +# assert torch.isfinite(torch.concat(list(activations.values()))).all() +# return activations diff --git a/src/extraction/config.py b/src/extraction/config.py index cffd50e..d501556 100644 --- a/src/extraction/config.py +++ b/src/extraction/config.py @@ -65,7 +65,7 @@ class ExtractConfig(Serializable): disable_ds_cache: bool = False """Disable huggingface datasets cache.""" - intervention_direction_method: str = "cluster_mean" + intervention_direction_method: str = "mm" """"how to intervent: pca, cluster_mean, random""" intervention_fit_examples: int = 200 diff --git a/src/models/load.py b/src/models/load.py index 21f430f..b12d0c4 100644 --- a/src/models/load.py +++ b/src/models/load.py @@ -58,8 +58,11 @@ def load_model(model_repo = "microsoft/phi-2", pad_token_id=0) -> Tuple[AutoMod model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options) - # from auto_gptq import exllama_set_max_input_length - # model = exllama_set_max_input_length(model, max_input_length=5000) + try: + from auto_gptq import exllama_set_max_input_length + model = exllama_set_max_input_length(model, max_input_length=5000) + except Exception as e: + logger.exception("could not set exllama max input length") return model, tokenizer diff --git a/src/repe/__init__.py b/src/repe/__init__.py index 589b175..5a37c16 100644 --- a/src/repe/__init__.py +++ b/src/repe/__init__.py @@ -4,8 +4,6 @@ warnings.filterwarnings("ignore") from .pipelines import repe_pipeline_registry -# RepReading -from .rep_readers import * from .rep_reading_pipeline import * # RepControl diff --git a/src/repe/interventions.py b/src/repe/interventions.py new file mode 100644 index 0000000..49b468b --- /dev/null +++ b/src/repe/interventions.py @@ -0,0 +1,226 @@ +import torch as t +from typing import Dict, List +from jaxtyping import Float +from torch import Tensor + +class Intervention(t.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x, iid=None): + raise NotImplementedError + + def pred(self, x, iid=None): + raise NotImplementedError + + def edit(self, x): + direction = self.direction + true_acts, false_acts = acts[labels==1], acts[labels==0] + true_mean, false_mean = true_acts.mean(0), false_acts.mean(0) + direction = direction / direction.norm() + diff = (true_mean - false_mean) @ direction + direction = diff * direction + + x[0][:,-1, :] += direction * alpha + return self(x) + + @staticmethod + def from_data(acts: Dict[int, Float[Tensor, "batch neurons"]], labels: List[bool], **kwargs) -> t.nn.Module: + """builds the class from data.""" + raise NotImplementedError + + +class LayerInterventions(t.nn.Module): + """An intervention for each layer""" + def __init__(self, interventions: dict): + super().__init__() + self.interventions = interventions + + def forward(self, x, **kwargs): + return {k: v(x[k], **kwargs) for k, v in self.interventions.items()} + + def pred(self, x, **kwargs): + return {k: v.pred(x[k], **kwargs) for k, v in self.interventions.items()} + + @staticmethod + def from_data(Intervention, acts: Dict[int, Float[Tensor, "batch neurons"]], labels: List[bool], layer_name_tmpl:str, **kwargs) -> t.nn.Module: + return LayerInterventions({layer_name_tmpl.format(layer_n): Intervention.from_data(act, labels, **kwargs) for layer_n, act in acts.items()}) + + @property + def direction(self): + return {k: v.direction for k, v in self.interventions.items()} + + +class MMProbe(t.nn.Module): + """ + Mean Mass Probe + + From geometry-of-truth repo + https://github.com/saprmarks/geometry-of-truth/blob/91b223224699754efe83bbd3cae04d434dda0760/probes.py#L35C1-L64C21 + """ + def __init__(self, direction, covariance=None, inv=None, atol=1e-3): + super().__init__() + self.direction = t.nn.Parameter(direction, requires_grad=False) + if inv is None: + self.inv = t.nn.Parameter(t.linalg.pinv(covariance, hermitian=True, atol=atol), requires_grad=False) + else: + self.inv = t.nn.Parameter(inv, requires_grad=False) + + def forward(self, x, iid=False): + self.to(x.device).to(x.dtype) + if iid: + return t.nn.Sigmoid()(x @ self.inv @ self.direction) + else: + return t.nn.Sigmoid()(x @ self.direction) + + def pred(self, x, iid=False): + return self(x, iid=iid).round() + + @staticmethod + def from_data(acts, labels, atol=1e-3, device='cpu'): + # acts: tensor of shape [n_activations, activation_dimension]. + pos_acts, neg_acts = acts[labels==1], acts[labels==0] + pos_mean, neg_mean = pos_acts.mean(0), neg_acts.mean(0) + direction = pos_mean - neg_mean + + centered_data = t.cat([pos_acts - pos_mean, neg_acts - neg_mean], 0) + covariance = centered_data.t() @ centered_data / acts.shape[0] + + probe = MMProbe(direction, covariance=covariance).to(device) + + return probe + + + +class COMProbe(t.nn.Module): + """ + Center of Mass Probe + + From honest llama repo redone as a probe class + https://github.com/likenneth/honest_llama/blob/207bb14b2c005e0593487cca8d22e072cbcb987b/utils.py#L730 + """ + + def __init__(self, direction, proj_val_std, alpha=15): + super().__init__() + self.direction = t.nn.Parameter(direction, requires_grad=False) + self.std = t.nn.Parameter(proj_val_std, requires_grad=False) + self.alpha = alpha + + def forward(self, x, iid=False): + # https://github.com/likenneth/honest_llama/blob/207bb14b2c005e0593487cca8d22e072cbcb987b/validation/validate_2fold.py#L116 + x += self.alpha * self.std * self.direction + return x + + def pred(self, x, iid=False): + return self(x, iid=iid).round() + + @staticmethod + def from_data(acts, labels, device='cpu'): + pos_acts, neg_acts = acts[labels==1], acts[labels==0] + pos_mean, neg_mean = pos_acts.mean(0), neg_acts.mean(0) + direction = pos_mean - neg_mean + direction = direction / t.linalg.norm(direction) + + proj_vals = acts @ direction.T + proj_val_std = t.std(proj_vals) + + probe = COMProbe(direction, proj_val_std=proj_val_std).to(device) + + return probe + + + +class LRProbe(t.nn.Module): + """ + Linear regression probe + From geometry-of-truth repo + """ + def __init__(self, d_in): + super().__init__() + self.net = t.nn.Sequential( + t.nn.Linear(d_in, 1, bias=False), + t.nn.Sigmoid() + ) + + def forward(self, x, iid=None): + return self.net(x).squeeze(-1) + + def pred(self, x, iid=None): + return self(x).round() + + @staticmethod + def from_data(acts, labels, lr=0.001, weight_decay=0.1, epochs=1000, device='cpu'): + acts, labels = acts.to(device), labels.to(device) + probe = LRProbe(acts.shape[-1]).to(device) + + opt = t.optim.AdamW(probe.parameters(), lr=lr, weight_decay=weight_decay) + for _ in range(epochs): + opt.zero_grad() + loss = t.nn.BCELoss()(probe(acts), labels) + loss.backward() + opt.step() + + return probe + + @property + def direction(self): + return self.net[0].weight.data[0] + + +def ccs_loss(probe, acts, neg_acts): + p_pos = probe(acts) + p_neg = probe(neg_acts) + consistency_losses = (p_pos - (1 - p_neg)) ** 2 + confidence_losses = t.min(t.stack((p_pos, p_neg), dim=-1), dim=-1).values ** 2 + return t.mean(consistency_losses + confidence_losses) + + +class CCSProbe(t.nn.Module): + """ + Contrast-Consistent Search + From geometry-of-truth repo + Originally from https://arxiv.org/pdf/2212.03827.pdf + """ + def __init__(self, d_in): + super().__init__() + self.net = t.nn.Sequential( + t.nn.Linear(d_in, 1, bias=False), + t.nn.Sigmoid() + ) + + def forward(self, x, iid=None): + return self.net(x).squeeze(-1) + + def pred(self, acts, iid=None): + return self(acts).round() + + @staticmethod + def from_data(acts, labels, lr=0.001, weight_decay=0.1, epochs=1000, device='cpu'): + pos_acts, neg_acts = acts[labels==1], acts[labels==0] + pos_acts, neg_acts = pos_acts.to(device), neg_acts.to(device) + probe = CCSProbe(pos_acts.shape[-1]).to(device) + + opt = t.optim.AdamW(probe.parameters(), lr=lr, weight_decay=weight_decay) + for _ in range(epochs): + opt.zero_grad() + loss = ccs_loss(probe, pos_acts, neg_acts) + loss.backward() + opt.step() + + if labels is not None: # flip direction if needed + acc = (probe.pred(pos_acts) == labels).float().mean() + if acc < 0.5: + probe.net[0].weight.data *= -1 + + return probe + + @property + def direction(self): + return self.net[0].weight.data[0] + +DIRECTION_FINDERS = { + 'ccs': CCSProbe, + 'lr': LRProbe, + 'com': COMProbe, + 'mm': MMProbe, +} diff --git a/src/repe/rep_control_pipeline_baukit.py b/src/repe/rep_control_pipeline_baukit.py index 6be86e9..eb41e88 100644 --- a/src/repe/rep_control_pipeline_baukit.py +++ b/src/repe/rep_control_pipeline_baukit.py @@ -1,3 +1,4 @@ +from hmac import new import re import torch from transformers.pipelines import ( @@ -13,10 +14,8 @@ from functools import partial from einops import rearrange from transformers.modeling_outputs import ModelOutput from src.datasets.scores import choice2ids, default_class2choices, logits2choice_probs2 -# from src.datasets.scores import scores2choice_probs from src.helpers.torch import clear_mem, detachcpu -from src.datasets.intervene import intervention_meta_fn2, Activations - +from src.repe.interventions import Intervention def try_half(v): @@ -35,29 +34,57 @@ def row_choice_ids(answer_choices, tokenizer): return choice2ids([c for c in answer_choices], tokenizer) -# def split_outputs(o): +def intervene(output, intervention) -> torch.Tensor: + """ + + https://github.com/saprmarks/geometry-of-truth/blob/91b223224699754efe83bbd3cae04d434dda0760/interventions.ipynb + """ + alpha = -1 + output[:, - 1, :] += intervention.direction * alpha + return output + +def intervention_fn(outputs: torch.Tensor, layer_name: str, intervention: Intervention) -> torch.Tensor: + """ + This adapts and intervention function for baukit Tracdict + + - 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: + """ + fn = intervention.interventions[layer_name] + + # different transformer have different formats of layer returns + if type(outputs) is tuple: + output0 = intervene(outputs[0], fn) + return (output0, *outputs[1:]) + elif type(outputs) is torch.Tensor: + return intervene(outputs, fn) + else: + raise ValueError(f"layer 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): + def __init__(self, model, tokenizer, max_length, **kwargs): super().__init__(model=model, tokenizer=tokenizer, **kwargs) self.max_length = max_length - self.layer_name_tmpl = layer_name_tmpl # self.default_class2choiceids = choice2ids(default_class2choices, tokenizer) def __call__(self, model_inputs, **kwargs): return super().__call__(model_inputs, **kwargs) - def _sanitize_parameters(self, activations=None, truncation=None, tokenize_kwargs=None, return_tensors=None, **kwargs): + def _sanitize_parameters(self, intervention=None, truncation=None, tokenize_kwargs=None, return_tensors=None, **kwargs): """This processed the init params.""" if tokenize_kwargs is None: tokenize_kwargs = {} preprocess_params = tokenize_kwargs - forward_params = {'activations': activations} + forward_params = {'intervention': intervention} postprocess_params = {} if return_tensors is not None: @@ -83,17 +110,12 @@ class RepControlPipeline2(FeatureExtractionPipeline): inputs["attention_mask"] = torch.tensor(inputs['attention_mask'], dtype=torch.bool, device=self.model.device) return inputs - def _forward(self, inputs: dict, activations: Dict[str, float]) -> ModelOutput: + def _forward(self, inputs: dict, intervention: Intervention) -> ModelOutput: assert inputs['input_ids'].ndim == 2, f"expected input_ids to be (batch, seq), got {inputs['input_ids'].shape}" # make intervention functions - layers_names = [self.layer_name_tmpl.format(i) for i in activations.keys()] - # FIXME: [0] is positive, [1] is negative. We can also multiply by -1, 0, or 1 - # FIXME clean this up, we are only using the first one, so it's confusing. either pass 1, or use both so the logic is in one place only - activations_pos_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations.items()}) - activations_neut = Activations({self.layer_name_tmpl.format(k):0. * v for k,v in activations.items()}) - edit_fn_pos = partial(intervention_meta_fn2, activations=activations_pos_i) - edit_fn_neu = partial(intervention_meta_fn2, activations=activations_neut) + layers_names = list(intervention.interventions.keys()) + edit_fn = partial(intervention_fn, intervention=intervention) self.model.eval() model_in = dict( @@ -119,14 +141,11 @@ class RepControlPipeline2(FeatureExtractionPipeline): # intervent in the negative and positive direction with torch.no_grad(): with TraceDict( - self.model, layers_names, detach=True, edit_output=edit_fn_pos + self.model, layers_names, detach=True, edit_output=edit_fn ) as ret: outputs_pos = transform_model_output(self.model(**model_in)) - with TraceDict( - self.model, layers_names, detach=True, edit_output=edit_fn_neu - ) as ret: - outputs_neg = transform_model_output(self.model(**model_in)) + outputs_neg = transform_model_output(self.model(**model_in)) # stack the outputs o = {k: torch.stack([outputs_neg[k], outputs_pos[k]], -1) for k in outputs_neg.keys()} diff --git a/src/repe/rep_readers.py b/src/repe/rep_readers.py deleted file mode 100644 index a9e0894..0000000 --- a/src/repe/rep_readers.py +++ /dev/null @@ -1,242 +0,0 @@ -from abc import ABC, abstractmethod -from sklearn.decomposition import PCA -from sklearn.cluster import KMeans -import numpy as np -from itertools import islice -import torch - -### 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) - direction = direction.to(H.device).squeeze(0) - mag = torch.linalg.norm(direction).to(H.device) - assert torch.isfinite(mag) - return (H[0] @ direction[:, None] # check that the dimensions match - 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] == 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[int(label)] else 0 for o, label in zip(projected_scores, train_choices)] - outputs_max = [1 if max(o) == o[int(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): - - # see also https://github.com/likenneth/honest_llama/blob/207bb14b2c005e0593487cca8d22e072cbcb987b/utils.py#L730 - - # 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 index cbb1b8c..1e0d2fa 100644 --- a/src/repe/rep_reading_pipeline.py +++ b/src/repe/rep_reading_pipeline.py @@ -2,7 +2,8 @@ from typing import List, Union, Optional from transformers import Pipeline import torch import numpy as np -from .rep_readers import DIRECTION_FINDERS, RepReader + +from .interventions import DIRECTION_FINDERS, Intervention, LayerInterventions class RepReadingPipeline(Pipeline): """Returns the directions for each layer, for each example.""" @@ -24,12 +25,12 @@ class RepReadingPipeline(Pipeline): 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() + hidden_states_layers[layer] = hidden_states.cpu().to(dtype=torch.float32).detach() return hidden_states_layers def _sanitize_parameters(self, - rep_reader: RepReader=None, + intervention: Intervention=None, rep_token: Union[str, int]=-1, hidden_layers: Union[List[int], int]=-1, component_index: int=0, @@ -45,8 +46,8 @@ class RepReadingPipeline(Pipeline): 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 + assert intervention is None or len(intervention.directions) == len(hidden_layers), f"expect total intervention directions ({len(intervention.directions)})== total hidden_layers ({len(hidden_layers)})" + forward_params['intervention'] = intervention forward_params['hidden_layers'] = hidden_layers forward_params['component_index'] = component_index forward_params['which_hidden_states'] = which_hidden_states @@ -65,7 +66,7 @@ class RepReadingPipeline(Pipeline): 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): + def _forward(self, model_inputs, rep_token, hidden_layers, intervention=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. @@ -80,38 +81,37 @@ class RepReadingPipeline(Pipeline): 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: + if intervention is None: return hidden_states - return rep_reader.transform(hidden_states, hidden_layers, component_index) + return intervention(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_layers=hidden_layers, batch_size=batch_size, intervention=None, which_hidden_states=which_hidden_states, **tokenizer_args) + # return hidden_states_outputs 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" + for layer in hidden_layers: + hidden_states[layer] = torch.stack(hidden_states[layer]) + return hidden_states def get_directions( self, train_inputs: Union[str, List[str], List[List[str]]], + train_labels: List[int], 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_method: str = 'mm', direction_finder_kwargs: dict = {}, which_hidden_states: Optional[str]=None, + layer_name_tmpl: str = "model.layers.{}", **tokenizer_args,): """Train a RepReader on the training data. Args: @@ -123,36 +123,21 @@ class RepReadingPipeline(Pipeline): 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) + Intervention = DIRECTION_FINDERS[direction_method] - # 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[1:]: - for _ in range(n_difference): - # FIXME: this is wrong, it's skipping batches... - relative_hidden_states[layer] = relative_hidden_states[layer] - relative_hidden_states[layer-1] - - # 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) + # 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) - return direction_finder + # get differences between pairs + relative_hidden_states = {k: torch.clone(v) for k, v in hidden_states.items()} + for layer in hidden_layers[1:]: + for _ in range(n_difference): + relative_hidden_states[layer] = relative_hidden_states[layer] - relative_hidden_states[layer-1] + + # fit probe + # TODO: use rel or abs? + probe = LayerInterventions.from_data(Intervention, relative_hidden_states, torch.LongTensor(train_labels), layer_name_tmpl) + + return probe