mirror of
https://github.com/wassname/discovering_latent_knowledge.git
synced 2026-09-09 11:21:22 +08:00
wip
This commit is contained in:
wassname
+38
-3
@@ -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 _
|
||||
|
||||
@@ -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": [
|
||||
"<src.repe.rep_readers.ClusterMeanRepReader at 0x7fdcf2279810>"
|
||||
]
|
||||
},
|
||||
"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": [
|
||||
"<src.repe.rep_readers.ClusterMeanRepReader at 0x7fdcf2279810>"
|
||||
]
|
||||
},
|
||||
"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": [
|
||||
"<src.repe.rep_control_pipeline_baukit.RepControlPipeline2 at 0x7fdbd7545a20>"
|
||||
]
|
||||
},
|
||||
"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 <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/notebooks/104_check_intervention.ipynb#Y102sZmlsZQ%3D%3D?line=0'>1</a>\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 <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/notebooks/104_check_intervention.ipynb#Y102sZmlsZQ%3D%3D?line=1'>2</a>\u001b[0m r1 \u001b[39m=\u001b[39m rep_control_pipeline2(model_inputs\u001b[39m=\u001b[39mds,\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/notebooks/104_check_intervention.ipynb#Y102sZmlsZQ%3D%3D?line=2'>3</a>\u001b[0m activations\u001b[39m=\u001b[39mactivations_neg_i,\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/notebooks/104_check_intervention.ipynb#Y102sZmlsZQ%3D%3D?line=3'>4</a>\u001b[0m batch_size\u001b[39m=\u001b[39mbatch_size,)\n\u001b[0;32m----> <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/notebooks/104_check_intervention.ipynb#Y102sZmlsZQ%3D%3D?line=4'>5</a>\u001b[0m r \u001b[39m=\u001b[39m \u001b[39mlist\u001b[39;49m(r1)\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/notebooks/104_check_intervention.ipynb#Y102sZmlsZQ%3D%3D?line=5'>6</a>\u001b[0m o \u001b[39m=\u001b[39m r[\u001b[39m0\u001b[39m]\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/notebooks/104_check_intervention.ipynb#Y102sZmlsZQ%3D%3D?line=6'>7</a>\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.<locals>.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__.<locals>.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",
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Generated
+1
-1
@@ -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"
|
||||
|
||||
+60
-60
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+5
-2
@@ -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
|
||||
|
||||
|
||||
@@ -4,8 +4,6 @@ warnings.filterwarnings("ignore")
|
||||
|
||||
from .pipelines import repe_pipeline_registry
|
||||
|
||||
# RepReading
|
||||
from .rep_readers import *
|
||||
from .rep_reading_pipeline import *
|
||||
|
||||
# RepControl
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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()}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user