tidy

mjc_notes.md

@@ -2253,5 +2253,20 @@ Some haven't worked along:
 - ranking (maybe because dropout doesnt give enougth variation? or the variation is useless), it needs intervention, but intervention might work on it's own. And then all ranking does it possibly help avoid overfitting!
 - intervention... it's hard to find a good one. It's either too little to matter or too much and the model is incoherent, meaning it's a implausible intervention. Plus all my interventions so far have been for the word true, not true or deception. 
   - [ ] Try doing MMProbe with deception vs truth?
+  - [ ] Try is residual not hidden...
   - [ ] Try SGD only on bias! trying to flip the probabilities on a large batch (can use grad accum)
-- SAE... no one has solved this. Maybe with an important matrix (which can come from an intervention)
+- SAE... no one has solved this. 
+  - [ ] Maybe with an important matrix (which can come from an intervention)
+
+
+also read a few
+- ognitive Dissonance: Why Do Language Model Outputs  https://arxiv.org/abs/2312.03729
+- https://lilianweng.github.io/posts/2018-08-12-vae/#vq-vae-and-vq-vae-2
+- geometry of truth https://arxiv.org/abs/2310.06824
+- openai weak to strong generalization
+
+
+Prioritise small experiments in notebooks
+- take all the recorded hidden states and seperate into truth and deception
+- try a normal intervention and test it
+- try novel sgb bias intervention

notebooks/105_new_interventions.ipynb

@@ -0,0 +1,273 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# A scratch pad to run model inference manually\n",
+    "\n",
+    "Prioritise small experiments in notebooks\n",
+    "- take all the recorded hidden states and seperate into truth and deception\n",
+    "- try a normal intervention and test it\n",
+    "- try novel sgb bias intervention"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "from matplotlib import pyplot as plt\n",
+    "\n",
+    "plt.style.use(\"ggplot\")\n",
+    "\n",
+    "from typing import Optional, List, Dict, Union\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch import Tensor\n",
+    "from torch import optim\n",
+    "from torch.utils.data import random_split, DataLoader, TensorDataset\n",
+    "\n",
+    "from pathlib import Path\n",
+    "import transformers\n",
+    "\n",
+    "\n",
+    "from loguru import logger\n",
+    "\n",
+    "logger.add(os.sys.stderr, format=\"{time} {level} {message}\", level=\"INFO\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load my code\n",
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "\n",
+    "from src.extraction.config import ExtractConfig\n",
+    "from src.prompts.prompt_loading import load_preproc_dataset\n",
+    "from src.models.load import load_model\n",
+    "from src.datasets.intervene import create_cache_interventions\n",
+    "from src.prompts.prompt_loading import load_prompt_structure\n",
+    "from src.repe import repe_pipeline_registry\n",
+    "\n",
+    "repe_pipeline_registry()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # config transformers\n",
+    "# from datasets import set_caching_enabled, disable_caching\n",
+    "# disable_caching()\n",
+    "# os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_name = \"phi-2-GPTQ_w_hidden_states\"\n",
+    "[str(s) for s in sorted(Path(\"../.ds/\").glob(f\"*{model_name}*\"))]\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_from_disk, concatenate_datasets\n",
+    "from src.datasets.load import ds2df, load_ds, get_ds_name\n",
+    "from src.datasets.load import ds2df, load_ds, get_ds_name, filter_ds_to_known\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "fs = [\n",
+    "    \"../.ds/wassname_phi-2-GPTQ_w_hidden_states_amazon_polarity_test_600\",\n",
+    "    # \"../.ds/wassname_phi-2-GPTQ_w_hidden_states_amazon_polarity_train_3600\",\n",
+    "    \"../.ds/wassname_phi-2-GPTQ_w_hidden_states_glue_qnli_test_600\",\n",
+    "    # \"../.ds/wassname_phi-2-GPTQ_w_hidden_states_glue_qnli_train_3600\",\n",
+    "    \"../.ds/wassname_phi-2-GPTQ_w_hidden_states_imdb_test_600\",\n",
+    "    # \"../.ds/wassname_phi-2-GPTQ_w_hidden_states_imdb_train_3600\",\n",
+    "    \"../.ds/wassname_phi-2-GPTQ_w_hidden_states_super_glue_boolq_test_600\",\n",
+    "    # \"../.ds/wassname_phi-2-GPTQ_w_hidden_states_super_glue_boolq_train_3600\",\n",
+    "]\n",
+    "dss = [load_ds(f) for f in fs]\n",
+    "dss\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from src.datasets.load import ds2df, load_ds, get_ds_name, filter_ds_to_known, qc_ds\n",
+    "for ds in dss:\n",
+    "    qc_ds(ds)\n",
+    "    # ds = ds.with_format(\"numpy\")\n",
+    "    \n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "select rows are 93.00% based on knowledge\n",
+      "select rows are 61.33% based on knowledge\n",
+      "select rows are 84.14% based on knowledge\n",
+      "select rows are 82.00% based on knowledge\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Dataset({\n",
+       "    features: ['end_hidden_states', 'end_logits', 'choice_probs', 'label_true', 'instructed_to_lie', 'question', 'answer_choices', 'choice_ids', 'template_name', 'sys_instr_name', 'example_i', 'input_truncated', 'truncated', 'text_ans', 'ans'],\n",
+       "    num_rows: 1870\n",
+       "})"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# combine\n",
+    "dss_known = [filter_ds_to_known(d) for d in dss]\n",
+    "ds = concatenate_datasets(dss_known)\n",
+    "ds = ds.with_format(\"numpy\")\n",
+    "ds\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "after filtering we have 137 num successful lies out of 1870 dataset rows\n"
+     ]
+    }
+   ],
+   "source": [
+    "# QC: make sure we didn't lose all of the successful lies, which would make the problem trivial\n",
+    "df2 = ds2df(ds)\n",
+    "df_subset_successull_lies = df2.query(\n",
+    "    \"instructed_to_lie==True & ((llm_ans==1)==label_instructed)\"\n",
+    ")\n",
+    "print(\n",
+    "    f\"after filtering we have {len(df_subset_successull_lies)} num successful lies out of {len(df2)} dataset rows\"\n",
+    ")\n",
+    "assert (\n",
+    "    len(df_subset_successull_lies) > 0\n",
+    "), \"there should be successful lies in the dataset\"\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_name = \"amazon_polarity\"\n",
+    "cfg = ExtractConfig(\n",
+    "    max_examples=(400, 400),\n",
+    "    intervention_fit_examples=160,\n",
+    ")\n",
+    "print(cfg)\n",
+    "batch_size = cfg.batch_size\n",
+    "\n",
+    "model, tokenizer = load_model(\n",
+    "    cfg.model, pad_token_id=cfg.pad_token_id, disable_exllama=False\n",
+    ")\n",
+    "print(model)\n",
+    "\n",
+    "N_train, N_test = cfg.max_examples\n",
+    "N = sum(cfg.max_examples)\n",
+    "ds_tokens = load_preproc_dataset(\n",
+    "    ds_name,\n",
+    "    tokenizer,\n",
+    "    N=N,\n",
+    "    seed=cfg.seed,\n",
+    "    num_shots=cfg.num_shots,\n",
+    "    max_length=cfg.max_length,\n",
+    "    prompt_format=cfg.prompt_format,\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Intervention"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

notebooks/make_dataset2.py

@@ -143,8 +143,6 @@ rep_control_pipeline2
 from src.datasets.intervene import test_intervention_quality
 
 
-
-
 # %%
 # test intervention quality
 # TODO perhaps move this to intervention create/load/cache
@@ -252,16 +250,24 @@ for ds_name in cfg.datasets:
     assert cfg.intervention_fit_examples < N_train
     N_train_split = N_train - cfg.intervention_fit_examples
 
-
     # split the dataset, it's preshuffled
     dataset_fit = ds_tokens.select(range(cfg.intervention_fit_examples))
-    dataset_train = ds_tokens.select(range(cfg.intervention_fit_examples, N_train_split))
+    dataset_train = ds_tokens.select(
+        range(cfg.intervention_fit_examples, N_train_split)
+    )
     dataset_test = ds_tokens.select(range(N_train_split, len(ds_tokens)))
     assert len(dataset_train) > 3, f"dataset_train is too small {len(dataset_train)}"
     assert len(dataset_test) > 3
 
     # FIXME:
-    test_intervention_quality(dataset_train, intervention, 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,

src/datasets/intervene.py

@@ -16,47 +16,6 @@ from loguru import logger
 from jaxtyping import Float
 from torch import nn, Tensor
 
-# Activations = NewType("Activations", Dict[str, torch.Tensor])
-
-# 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, :]
-
-#     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
-
-#     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(
     model,
@@ -74,9 +33,9 @@ def create_cache_interventions(
 
     So we always load a cached version if possible. to make it approx repeatable use the same dataset etc
     """
-    direction_method=cfg.intervention_direction_method
-    N_fit_examples=cfg.intervention_fit_examples
-    batch_size=cfg.batch_size
+    direction_method = cfg.intervention_direction_method
+    N_fit_examples = cfg.intervention_fit_examples
+    batch_size = cfg.batch_size
     tokenizer_args = dict(
         padding="max_length",
         max_length=cfg.max_length,
@@ -126,9 +85,7 @@ def create_cache_interventions(
             **tokenizer_args,
         )
 
-        assert np.isfinite(
-            np.concatenate(list(intervention.direction.values()))
-        ).all()
+        assert np.isfinite(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:
@@ -142,7 +99,6 @@ def create_cache_interventions(
     return intervention
 
 
-
 def test_intervention_quality(
     dataset_train, intervention, model, rep_control_pipeline2, batch_size=2, ds_name=""
 ):
@@ -150,7 +106,7 @@ def test_intervention_quality(
     Check interventions are ordered and different and valid
     """
     # TODO over multiple batches?
-    inputs = dataset_train#[:batch_size]
+    inputs = dataset_train  # [:batch_size]
     model.eval()
     baseline_outputs = []
     for batch_index in range(len(inputs) // batch_size):
@@ -176,7 +132,7 @@ def test_intervention_quality(
 
         # TODO coverage too
         # mean_prob = np.sum(r['choice_probs'], 1).mean()
-        coverage.append(torch.sum(r['choice_probs'], 1))
+        coverage.append(torch.sum(r["choice_probs"], 1))
 
         choice_true = choices[label]
         if label == 0:
@@ -191,11 +147,11 @@ def test_intervention_quality(
                 label=label,
             )
         )
-        if bi<5:
+        if bi < 5:
             print(f"\t Score: {list(ans)} of true ans `{choice_true}`=={label}")
             print(f"\t Ordered? {ordered} {np.argsort(ans)}")
             print(f"\t Different? {abs(np.diff(ans))>0.1} {np.diff(ans)}")
-            print("'\t top choices", r['text_ans'], 'should be valid')
+            print("'\t top choices", r["text_ans"], "should be valid")
     df = pd.DataFrame(data)
     print(f"N={len(df)}")
     print(f"rows that are ordered {df['order'].mean():%}")
@@ -207,22 +163,62 @@ 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"""
+from IPython.display import display, HTML
 
-#     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
+def top_toke_probs(
+    o,
+    tokenizer,
+    N=20,
+):
+    """
+        nicely return top token probabilities e.g.
+
+            prob_0	tokens_0	id_0	prob_1	tokens_1	id_1
+    0	0.811495	`Neg`	32863	0.666955	`Neg`	32863
+    1	0.113310	`Pos`	21604	0.194095	`Pos`	21604
+    2	0.020635	`Ne`	8199	0.065014	`Ne`	8199
+    """
+    data = {}
+    for i in range(o["end_logits"].shape[1]):
+        probs = torch.softmax(o["end_logits"][:, i], -1)
+        top = probs.argsort(0, descending=True)
+        top_probs = probs[top]
+        tokens_top20 = tokenizer.batch_decode(
+            top[:N], skip_special_tokens=False, clean_up_tokenization_spaces=False
+        )
+        tokens_top20 = [f"`{t}`" for t in tokens_top20]
+        data.update(
+            {
+                f"prob_{i}": top_probs[:N],
+                f"tokens_{i}": tokens_top20,
+                f"id_{i}": top[:N],
+            }
+        )
+    return pd.DataFrame(data)
+
+
+def print_pipeline_row(o: dict, tokenizer):
+    """take in single pipeline output, and prince intervention metrics."""
+    choices = [tokenizer.batch_decode(cc) for cc in o["choice_ids"]]
+    index = [o[0] for o in choices]
+    d = pd.DataFrame(
+        o["choice_probs"].numpy(), columns=["edit=None", "edit=+"], index=index
+    ).T
+    # d["top1 coverage"] = d.sum(1)
+    mean_prob = o["choice_probs"].sum(0)
+    d["coverage"] = mean_prob
+
+    print("choices", choices)
+    max_prob, max_token = torch.softmax(o["end_logits"][:, :], 0).max(0)
+    max_detoken = tokenizer.batch_decode(max_token)
+    d["top_token"] = max_detoken
+    d["top_prob"] = max_prob
+    d["label_true"] = o["label_true"]
+    d["label_instructed"] = o["label_instructed"]
+    print("choice probs")
+    display(d)
+
+    d1 = top_toke_probs(o, tokenizer)
+    print("top token probs")
+    display(d1)