cache intervention, fixing mem leak

mjc_notes.md

@@ -1828,4 +1828,16 @@ TODO:
 
 OK so I got my pipeline working, and it's the same. And it uses the PCA.
 
-Next step simply, and put into a dataset. Then I can try the ranking
+
+# 2023-10-24 16:58:30
+
+ - [x] Next step simply, 
+ - [ ] and put into a dataset. 
+   - [ ] make sure I clean up the dataset script
+ - [ ] Then I can try the ranking
+
+
+we need to save cache load
+honesty_rep_reader both directions and signs!
+or activations
+

notebooks/201_make_data.ipynb

@@ -73,34 +73,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "ExtractConfig(datasets=('amazon_polarity', 'super_glue:boolq', 'glue:qnli', 'imdb'), model='TheBloke/WizardCoder-Python-13B-V1.0-GPTQ', data_dirs=(), max_examples=(100, 100), num_shots=1, num_variants=-1, layers=(), seed=42, template_path=None, max_length=555)\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[32m2023-10-24 16:57:59.555\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging pad_token_id from 32000 to 0\u001b[0m\n",
-      "2023-10-24T16:57:59.555939+0800 INFO changing pad_token_id from 32000 to 0\n",
-      "\u001b[32m2023-10-24 16:57:59.557\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging padding_side from right to left\u001b[0m\n",
-      "2023-10-24T16:57:59.557026+0800 INFO changing padding_side from right to left\n",
-      "\u001b[32m2023-10-24 16:57:59.557\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging truncation_side from right to left\u001b[0m\n",
-      "2023-10-24T16:57:59.557571+0800 INFO changing truncation_side from right to left\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# model_name_or_path = \"TheBloke/Wizard-Vicuna-30B-Uncensored-GPTQ\"\n",
     "# model_name_or_path = \"TheBloke/Mistral-7B-Instruct-v0.1-GPTQ\"\n",
     "model_name_or_path = \"TheBloke/WizardCoder-Python-13B-V1.0-GPTQ\"\n",
     "\n",
+    "batch_size = 2\n",
+    "\n",
     "cfg = ExtractConfig(max_examples=(100, 100), model=model_name_or_path)\n",
     "print(cfg)\n",
     "\n",
@@ -113,20 +95,110 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "rep_token = -1\n",
-    "batch_size = 2\n",
+    "\n",
+    "\n",
     "# hidden_layers = list(range(-1, -model.config.num_hidden_layers, -1))\n",
     "# hidden_layers = [f\"model.layers.{i}\" for i in range(8, model.config.num_hidden_layers, 3)]\n",
-    "hidden_layers = list(range(8, model.config.num_hidden_layers, 3))\n",
-    "hidden_layers           \n",
+    "# hidden_layers = list(range(8, model.config.num_hidden_layers, 3))\n",
+    "# hidden_layers\n",
     "\n",
-    "n_difference = 1\n",
-    "direction_method = 'pca'\n",
-    "rep_reading_pipeline =  pipeline(\"rep-reading\", model=model, tokenizer=tokenizer)\n",
-    "rep_reading_pipeline\n",
+    "\n",
+    "\n",
+    "# rep_reading_pipeline =  pipeline(\"rep-reading\", model=model, tokenizer=tokenizer)\n",
+    "# rep_reading_pipeline\n",
+    "# hidden_layers\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # cache busting for the transformers map and ds steps\n",
+    "# !rm -rf ~/.cache/huggingface/datasets/generator\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Intervention fit/load"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pickle\n",
+    "N_fit_examples = 10\n",
+    "from src.config import root_folder\n",
+    "tokenizer_args=dict(padding=\"max_length\", max_length=cfg.max_length, truncation=True, add_special_tokens=True)\n",
+    "            \n",
+    "def load_rep_reader(model, tokenizer, cfg, N_fit_examples=20, batch_size=2):\n",
+    "    \"\"\"\n",
+    "    We want one set of interventions per model\n",
+    "    \n",
+    "    So we always load a cached version if possible. to make it approx repeatable use the same dataset etc\n",
+    "    \"\"\"\n",
+    "    model_name = cfg.model.replace('/', '-')\n",
+    "    intervention_f = root_folder / 'data' / 'interventions' / f'{model_name}.pkl'\n",
+    "    intervention_f.parent.mkdir(exist_ok=True, parents=True)\n",
+    "    if not intervention_f.exists():\n",
+    "        rep_token = -1\n",
+    "        n_difference = 1\n",
+    "        direction_method = 'pca'\n",
+    "        \n",
+    "        hidden_layers = list(range(8, model.config.num_hidden_layers, 3))\n",
+    "        \n",
+    "        dataset_fit = load_preproc_dataset('imdb', cfg, tokenizer, N=N_fit_examples)\n",
+    "        \n",
+    "        rep_reading_pipeline =  pipeline(\"rep-reading\", model=model, tokenizer=tokenizer)\n",
+    "        honesty_rep_reader = rep_reading_pipeline.get_directions(\n",
+    "            dataset_fit['question'], \n",
+    "            rep_token=rep_token, \n",
+    "            hidden_layers=hidden_layers, \n",
+    "            n_difference=n_difference, \n",
+    "            train_labels=dataset_fit['label_true'], \n",
+    "            direction_method=direction_method,\n",
+    "            batch_size=batch_size,\n",
+    "            **tokenizer_args\n",
+    "        )\n",
+    "        # and save\n",
+    "        with open(intervention_f, 'wb') as f:\n",
+    "            pickle.dump(honesty_rep_reader, f)\n",
+    "            logger.info(f'Saved interventions to {intervention_f}')\n",
+    "    else:\n",
+    "        with open(intervention_f, 'rb') as f:\n",
+    "            honesty_rep_reader = pickle.load(f)\n",
+    "        logger.info(f'Loaded interventions from {intervention_f}')\n",
+    "            \n",
+    "    return honesty_rep_reader\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "N_fit_examples = 20\n",
+    "honesty_rep_reader = load_rep_reader(model, tokenizer, cfg, N_fit_examples=N_fit_examples, batch_size=batch_size)\n",
+    "\n",
+    "hidden_layers = honesty_rep_reader.directions.keys()\n",
     "hidden_layers\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Dataset"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -136,16 +208,9 @@
     "# load dataset\n",
     "ds_name = 'imdb'\n",
     "ds_tokens = load_preproc_dataset(ds_name, cfg, tokenizer)\n",
-    "ds_tokens\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "N_fit_examples = 10\n",
+    "ds_tokens\n",
+    "\n",
+    "\n",
     "N_train_split = (len(ds_tokens) - N_fit_examples) //2\n",
     "\n",
     "# split the dataset, it's preshuffled\n",
@@ -161,28 +226,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tokenizer_args=dict(padding=\"max_length\", max_length=cfg.max_length, truncation=True, add_special_tokens=True)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# fit\n",
-    "train_labels = dataset_fit['label_true']\n",
-    "honesty_rep_reader = rep_reading_pipeline.get_directions(\n",
-    "    dataset_fit['question'], \n",
-    "    rep_token=rep_token, \n",
-    "    hidden_layers=hidden_layers, \n",
-    "    n_difference=n_difference, \n",
-    "    train_labels=dataset_fit['label_true'], \n",
-    "    direction_method=direction_method,\n",
-    "    batch_size=batch_size,\n",
-    "    **tokenizer_args\n",
-    ")\n",
-    "honesty_rep_reader\n"
+    "# # fit\n",
+    "# # FIXME: load or save here if the inputs are the same....\n",
+    "# train_labels = dataset_fit['label_true']\n",
+    "# honesty_rep_reader = rep_reading_pipeline.get_directions(\n",
+    "#     dataset_fit['question'], \n",
+    "#     rep_token=rep_token, \n",
+    "#     hidden_layers=hidden_layers, \n",
+    "#     n_difference=n_difference, \n",
+    "#     train_labels=dataset_fit['label_true'], \n",
+    "#     direction_method=direction_method,\n",
+    "#     batch_size=batch_size,\n",
+    "#     **tokenizer_args\n",
+    "# )\n",
+    "# honesty_rep_reader\n"
    ]
   },
   {
@@ -214,7 +271,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "inputs = dataset_train.select([0, 10, 30])[:3]\n",
+    "inputs = dataset_train[:3]\n",
     "inputs['question']\n"
    ]
   },
@@ -265,11 +322,18 @@
     "    \"rep-control2\", \n",
     "    model=model, \n",
     "    tokenizer=tokenizer, \n",
-    "    layers=layer_id, \n",
+    "    layers=hidden_layers, \n",
     "    max_length=cfg.max_length,)\n",
     "rep_control_pipeline2\n"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -278,13 +342,10 @@
    "source": [
     "\n",
     "coeff=8.0\n",
-    "max_new_tokens=3\n",
-    "text_gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens, use_cache=False, \n",
-    "                       output_hidden_states=True, return_dict=True, max_length=cfg.max_length,\n",
-    "                       )\n",
+    "max_new_tokens=1\n",
     "\n",
     "activations = {}\n",
-    "for layer in layer_id:\n",
+    "for layer in hidden_layers:\n",
     "    activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer]).to(model.device).half()\n",
     "    \n",
     "\n",
@@ -292,13 +353,149 @@
     "\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
-    "    baseline_outputs = rep_control_pipeline2(inputs, batch_size=batch_size, **text_gen_kwargs)\n",
-    "    control_outputs = rep_control_pipeline2(inputs, activations=activations, batch_size=batch_size,  **text_gen_kwargs)\n",
-    "    control_outputs_neg = rep_control_pipeline2(inputs, activations=activations_neg, batch_size=batch_size, **text_gen_kwargs)\n",
+    "    baseline_outputs = rep_control_pipeline2(inputs, batch_size=batch_size)\n",
+    "    control_outputs = rep_control_pipeline2(inputs, activations=activations, batch_size=batch_size)\n",
+    "    control_outputs_neg = rep_control_pipeline2(inputs, activations=activations_neg, batch_size=batch_size)\n",
     "\n",
     "\n",
     "metrics(control_outputs_neg, baseline_outputs, control_outputs)\n"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torch.utils.data import Dataset\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# TODO we need to save and cache for many split and datasets\n",
+    "rep_control_pipeline2\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "# from datasets import Dataset, DatasetInfo\n",
+    "import datasets\n",
+    "from src.config import root_folder\n",
+    "from pathvalidate import sanitize_filename\n",
+    "from src.helpers.ds import ds_keep_cols\n",
+    "\n",
+    "\n",
+    "def create_hs_ds(ds_name, ds_tokens, pipeline, activations=None, f = None, batch_size=2, split_type=\"train\"):\n",
+    "    \"create a dataset of hidden states.\"\"\"\n",
+    "    \n",
+    "    N = len(ds_tokens)\n",
+    "    dataset_name = sanitize_filename(f\"{cfg.model}_{ds_name}_{split_type}_{N}\", replacement_text=\"_\")\n",
+    "    f = root_folder / '.ds'/ f\"{dataset_name}\"\n",
+    "    \n",
+    "    info_kwargs = dict(extract_cfg=cfg.to_dict(), ds_name=ds_name, split_type=split_type, f=f, date=pd.Timestamp.now().isoformat(),)\n",
+    "    \n",
+    "    torch_cols = ['input_ids', 'attention_mask', 'choice_ids', 'question', 'answer_choices', 'example_i', 'label_true', 'sys_instr_name', 'template_name', 'instructed_to_lie']\n",
+    "    ds_t_subset = ds_keep_cols(ds_tokens, torch_cols)\n",
+    "    ds = ds_t_subset.to_iterable_dataset()\n",
+    "    # pipeline_it = rep_control_pipeline2(ds, batch_size=batch_size, **text_gen_kwargs)\n",
+    "    \n",
+    "    # first we make the calibration dataset with no intervention\n",
+    "    gen_kwargs = dict(\n",
+    "        model_inputs=ds,\n",
+    "        activations=activations,\n",
+    "        batch_size=batch_size,\n",
+    "    )\n",
+    "    \n",
+    "    ds1 = datasets.Dataset.from_generator(\n",
+    "        generator=pipeline,\n",
+    "        info=datasets.DatasetInfo(\n",
+    "            description=json.dumps(info_kwargs, indent=2),\n",
+    "            config_name=f,\n",
+    "        ),\n",
+    "        gen_kwargs=gen_kwargs,\n",
+    "        num_proc=1,\n",
+    "    )\n",
+    "    return ds1\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "ds1 = create_hs_ds('imdb', dataset_train, rep_control_pipeline2, split_type=\"train\")\n",
+    "ds1\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset_train['answer_choices'][0]\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torch.utils.data import Dataset\n",
+    "ds = dataset_train.to_iterable_dataset()\n",
+    "next(iter(ds))['answer_choices']\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# To Datasets\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "baseline_outputs\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

notebooks/make_dataset.py

@@ -24,7 +24,7 @@ from torch import Tensor
 import pickle
 import hashlib
 from pathlib import Path
-
+from pathvalidate import sanitize_filename
 import transformers
 from transformers import GPTQConfig
 from datasets import Dataset, DatasetInfo
@@ -278,7 +278,7 @@ def load_preproc_dataset(ds_name: str, cfg: ExtractConfig, tokenizer: PreTrained
 
 
 def row_choice_ids(r, tokenizer):
-    return choice2ids([[c] for c in r['answer_choices']], tokenizer)
+    return choice2ids([c for c in r['answer_choices']], tokenizer)
 
 
 def expand_choices(choices: List[str]) -> Set[str]:
@@ -419,11 +419,8 @@ if __name__ == "__main__":
 
     model, tokenizer = load_model(cfg.model)
 
-    
-
-    sanitize = lambda s:s.replace('/', '').replace('-', '_') if s is not None else s
     ds_name = 'imdb'
-    model_name = sanitize(cfg.model)
+    model_name = sanitize_filename(cfg.model)
     intervention, intervention_fn = load_intervention(ds_name, cfg, model, tokenizer, model_name)
 
     for ds_name in ds_names:
@@ -433,7 +430,7 @@ if __name__ == "__main__":
         # ## Save as Huggingface Dataset
         # get dataset filename
         N = len(ds_tokens)
-        dataset_name = f"{sanitize(cfg.model)}_{ds_name}_{split_type}_{N}"
+        dataset_name = f"{sanitize_filename(cfg.model)}_{ds_name}_{split_type}_{N}"
         f = root_folder / '.ds'/ f"{dataset_name}"
         
         ds1 = create_hs_ds(ds_name, ds_tokens, model, cfg, intervention_dicts=intervention, f=str(f))

poetry.lock

@@ -2001,6 +2001,21 @@ files = [
     {file = "pathspec-0.11.2.tar.gz", hash = "sha256:e0d8d0ac2f12da61956eb2306b69f9469b42f4deb0f3cb6ed47b9cce9996ced3"},
 ]
 
+[[package]]
+name = "pathvalidate"
+version = "3.2.0"
+description = "pathvalidate is a Python library to sanitize/validate a string such as filenames/file-paths/etc."
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "pathvalidate-3.2.0-py3-none-any.whl", hash = "sha256:cc593caa6299b22b37f228148257997e2fa850eea2daf7e4cc9205cef6908dee"},
+    {file = "pathvalidate-3.2.0.tar.gz", hash = "sha256:5e8378cf6712bff67fbe7a8307d99fa8c1a0cb28aa477056f8fc374f0dff24ad"},
+]
+
+[package.extras]
+docs = ["Sphinx (>=2.4)", "sphinx-rtd-theme (>=1.2.2)", "urllib3 (<2)"]
+test = ["Faker (>=1.0.8)", "allpairspy (>=2)", "click (>=6.2)", "pytest (>=6.0.1)", "pytest-discord (>=0.1.4)", "pytest-md-report (>=0.4.1)"]
+
 [[package]]
 name = "peft"
 version = "0.5.0"
@@ -3868,4 +3883,4 @@ multidict = ">=4.0"
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.10,<3.13"
-content-hash = "6068bb1c5ec5042e14fe98ec08112adb751f51829405378a393df028baf8c881"
+content-hash = "b8ab01ec55c8189d322df8be511d43175149fdbc5d5f11c91e7a5f1d91748f24"

pyproject.toml

@@ -26,6 +26,7 @@ baukit = {git = "https://github.com/davidbau/baukit.git"}
 eleuther-elk = "0.1.1"
 scikit-learn = "^1.3.1"
 pytorch-optimizer = "^2.12.0"
+pathvalidate = "^3.2.0"
 
 [[tool.poetry.source]]
 name = "pytorch"

src/datasets/hs.py

@@ -28,7 +28,7 @@ import torch.nn.functional as F
 from baukit.nethook import Trace, TraceDict, recursive_copy
 from einops import rearrange, reduce, repeat
 from src.datasets.scores import choice2id, choice2ids
-from src.helpers.torch import clear_mem
+from src.helpers.torch import clear_mem, detachcpu
 from collections import defaultdict
 from dataclasses import field
 from src.datasets.intervene import InterventionDict, intervention_meta_fn
@@ -217,15 +217,4 @@ class ExtractHiddenStates:
         layers_inds = sorted(set(list(strided_layers)+list(last_few)))
         return [layer_names[i] for i in layers_inds]
 
-def detachcpu(x):
-    """
-    Trys to convert torch if possible a single item
-    """
-    if isinstance(x, torch.Tensor):
-        # note apache parquet doesn't support half to we go for float https://github.com/huggingface/datasets/issues/4981
-        x = x.detach().cpu().float()
-        if x.squeeze().dim()==0:
-            return x.item()
-        return x
-    else:
-        return x
+

src/datasets/scores.py

@@ -10,15 +10,8 @@ from transformers import (
     PreTrainedModel
 )
 
-default_class2choices = {False: ['No', 'Negative', 'no', 'false', 'wrong', 'False'], True: ['Yes', 'Positive', 'yes', 'true', 'correct', 'right', 'True']}
+default_class2choices = [['No', 'Negative', 'negative', 'no', 'false', 'wrong', 'False', '0'], ['Yes', 'Positive', 'positive', 'yes', 'true', 'correct', 'right', 'True', '1']]
 
-# def class2choices_to_choices(class2choices):
-#     return [class2choices[i][0] for i in sorted(class2choices)]
-
-# def label_to_choice(label: bool, class2choices=default_class2choices) -> str:
-#     """turns a label like 0 to a choice like No"""
-#     choices = class2choices_to_choices(class2choices)
-#     return choices[label]
 
 def scores2choice_probs(row, class2_ids: List[List[int]], keys=["scores0"], prefix=""):
     """ Given next_token scores (logits) we take only the subset the corresponds to our
@@ -79,23 +72,14 @@ def choice2id(tokenizer, c: str, whitespace_first=False) -> List[int]:
 
 def choice2ids(all_choices: List[List[str]], tokenizer: PreTrainedTokenizer) -> List[List[int]]:
     choices = [list(itertools.chain(*[choice2id(tokenizer, c) for c in choices])) for choices in all_choices]
-    assert choices[0]!=choices[1], "choices should be different"
+    assert choices[0]!=choices[1], f"choices should be different but were not {all_choices}"
     assert choices[0][0]!=choices[1][0], "choices should be different"
     return choices
 
-# def get_choice_as_token(tokenizer, choice: str) -> int:
-#     return get_choices_as_tokens(tokenizer, [choice])[0]
 
-# def get_choices_as_tokens(
-#     tokenizer, choices:List[str] = ["Positive"], whitespace_first=True
-# ) -> List[int]:
-    
-    
-#     ids = []
-#     for c in choices:
-#         try:
-#             id_ = choice2id(tokenizer, c)
-#             ids.append(id_)
-#         except AssertionError as e:
-#             print(e)
-#     return ids
+def scores2choice_probs2(logits, choiceids: List[List[int]]):
+    """calculate the probability for each group of choices."""
+    assert logits.ndim==1, f"expected logits to be 1d, got {logits.shape}"
+    probs = torch.softmax(logits, 0)  # shape [tokens, inferences)
+    probs_c = torch.tensor([[probs[cc] for cc in c] for c in choiceids]).sum(1)  # sum over alternate choices e.g. [['decrease', 'dec'],['inc', 'increase']]
+    return probs_c

src/helpers/__init__.py

@@ -5,3 +5,4 @@ def bool2switch(x):
 def switch2bool(x):
     """[-1,1]->[0,1]"""
     return (x+1)/2
+

src/helpers/torch.py

@@ -44,3 +44,16 @@ def clear_mem():
     gc.collect()
     torch.cuda.empty_cache()
     gc.collect()
+
+def detachcpu(x):
+    """
+    Trys to convert torch if possible a single item
+    """
+    if isinstance(x, torch.Tensor):
+        # note apache parquet doesn't support half to we go for float https://github.com/huggingface/datasets/issues/4981
+        x = x.detach().cpu().float()
+        if x.squeeze().dim()==0:
+            return x.item()
+        return x
+    else:
+        return x

src/prompts/prompt_loading.py

@@ -232,6 +232,12 @@ def _convert_to_prompts(
 
     for template in templates:
         answer_choices=template.get_fixed_answer_choices_list()
+        
+        # skip prompts where the responses are similar in the first token
+        if answer_choices[0][:3]==answer_choices[1][:3]:
+            logger.debug(f"skipping prompt because it's answers are not unique (for the first token): {template.name} {answer_choices}")
+            continue
+        answer_choices = [[c] for c in answer_choices]
         for instructed_to_lie in [False, True]:
             for sys_instr_name, sys_instr in sys_instructions[instructed_to_lie].items():
                 fake_example = example.copy()
@@ -256,7 +262,7 @@ def _convert_to_prompts(
                     ]
                     for d in fewshot_texts:
                         # some of the answers have extra trailing text, that's OK. But extra preceeding text is not, let's check for that
-                        assert any([d['response'].startswith(a) for a in answer_choices]), f"fewshot response `{d['response']}` has extra preceeding text compared to allowed choices: {answer_choices}. template is: {template.name}"
+                        assert any([any([d['response'].startswith(a) for a in ac]) for ac in answer_choices]), f"fewshot response `{d['response']}` has extra preceeding text compared to allowed choices: {answer_choices}. template is: {template.name}"
                     prompt_parts = fewshot_texts + prompt_parts
                     
                 prompt_parts[0]['system'] = sys_instr

src/repe/rep_control_pipeline_baukit.py

@@ -1,3 +1,4 @@
+import re
 import torch
 from transformers.pipelines import (
     TextGenerationPipeline,
@@ -5,17 +6,26 @@ from transformers.pipelines import (
     Pipeline,
 )
 from transformers.pipelines.base import GenericTensor
+from datasets import Dataset
 from typing import List, Tuple, Dict, Any, Union, NewType
 from baukit.nethook import Trace, TraceDict, recursive_copy
 from functools import partial
-from src.datasets.scores import choice2ids
+from src.datasets.scores import choice2ids, default_class2choices, scores2choice_probs2
 # from src.datasets.scores import scores2choice_probs
+from src.helpers.torch import clear_mem, detachcpu
+
+Activations = NewType("Activations", Dict[str, torch.Tensor])
 
 
-Activations = NewType("InterventionDict", Dict[str, torch.Tensor])
+def hacky_sanitize_outputs(o):
+    """I can't find the mem leak, so lets just detach, cpu, clone, freemem."""
+    o = {k: detachcpu(v) for k, v in o.items()}
+    o = recursive_copy(o, detach=True, clone=True)
+    clear_mem()
+    return o
 
 def row_choice_ids(answer_choices, tokenizer):
-    return choice2ids([[c] for c in answer_choices], tokenizer)
+    return choice2ids([c for c in answer_choices], tokenizer)
 
 
 def intervene(output, activation):
@@ -51,6 +61,8 @@ class RepControlPipeline2(FeatureExtractionPipeline):
         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, activations=None, **kwargs):
         if activations is not None:
@@ -65,24 +77,30 @@ class RepControlPipeline2(FeatureExtractionPipeline):
             outputs = super().__call__(model_inputs, **kwargs)
         return outputs
     
-    def preprocess(self, inputs, **tokenize_kwargs) -> Dict[str, GenericTensor]:
+    def preprocess(self, inputs: Dataset, **tokenize_kwargs) -> Dict[str, GenericTensor]:
         # tokenize a batch of inputs
         return_tensors = self.framework
-        model_inputs = self.tokenizer(inputs['question'], return_tensors=return_tensors, return_attention_mask=True, add_special_tokens=True, truncation=True, padding="max_length", max_length=self.max_length, **tokenize_kwargs)
-        return {**inputs, **model_inputs}
+        if 'input_ids' not in inputs:
+            model_inputs = self.tokenizer(inputs['question'], return_tensors=return_tensors, return_attention_mask=True, add_special_tokens=True, truncation=True, padding="max_length", max_length=self.max_length, **tokenize_kwargs)
+            return {**inputs, **model_inputs}
+        else:
+            return inputs
 
     def _forward(self, model_inputs):
         inputs = dict(input_ids=model_inputs['input_ids'], attention_mask=model_inputs['attention_mask'])
         inputs.update(
             {"use_cache": False, "output_hidden_states": True, "return_dict": True}
         )
+        self.model.eval()
         with torch.no_grad():
             model_outputs = self.model(**inputs)
+        # o = {k: detachcpu(v) for k, v in o.items()}
+        # o = recursive_copy(o)
+        # clear_mem()
         
         # retain some of the inputs
-        keep_cols = ["answer_choices", "input_ids", "attention_mask"]
         model_outputs = {**model_inputs, **model_outputs}
-        return model_outputs
+        return hacky_sanitize_outputs(model_outputs)
 
     def postprocess(self, o):
         # note this sometimes deals with a batch, sometimes with a single result. infuriating
@@ -90,20 +108,22 @@ class RepControlPipeline2(FeatureExtractionPipeline):
         # This is called once for each result, but the text pipeline is set up to hande multiple...
         # This is called once for each result, but the text pipeline is set up to hande multiple...
         o["end_logits"] = o["logits"][:, -1, :].float()
-        # hidden_states = list(o.hidden_states)
         o["input_truncated"] = self.tokenizer.batch_decode(o['input_ids'])
         o["truncated"] = torch.sum(o["attention_mask"], 1)==self.max_length
         o["text_ans"] = self.tokenizer.batch_decode(o["end_logits"].argmax(-1))
-        o['choice_ids'] = [row_choice_ids(ac, self.tokenizer) for ac in o['answer_choices']] 
+        
+        if 'answer_choices' in o:
+            answer_choices = o['answer_choices']
+            if isinstance(answer_choices[0][0], str):
+                answer_choices = [answer_choices]
+        else:
+            answer_choices = default_class2choices # self.default_class2choiceids
+        
+        o['choice_ids'] = [row_choice_ids(ac, self.tokenizer) for ac in answer_choices]
         
         p = o['add_ans'] = torch.stack([scores2choice_probs2(l, c) for l, c in zip(o['end_logits'], o['choice_ids'])])
         o['ans'] = p[:, 1] / (torch.sum(p, 1) + 1e-5)
         
+        o = hacky_sanitize_outputs(o)
         return o
 
-def scores2choice_probs2(logits, choiceids: List[List[int]]):
-    """calculate the probability for each group of choices."""
-    assert logits.ndim==1, f"expected logits to be 1d, got {logits.shape}"
-    probs = torch.softmax(logits, 0)  # shape [tokens, inferences)
-    probs_c = torch.tensor([[probs[cc] for cc in c] for c in choiceids]).sum(1)  # sum over alternate choices e.g. [['decrease', 'dec'],['inc', 'increase']]
-    return probs_c