misc

mjc_notes.md

@@ -1841,3 +1841,9 @@ we need to save cache load
 honesty_rep_reader both directions and signs!
 or activations
 
+
+
+bugs:
+- [ ] why is it sometimes tensors sometimes not
+- [ ] why is it sometimes batched sometimes not
+- [ ] I want to force it into batched preproc, batched forward->postproce yeild single outputs

notebooks/201_make_data.ipynb

@@ -64,7 +64,7 @@
     "\n",
     "from src.models.load import load_model\n",
     "from src.extraction.config import ExtractConfig\n",
-    "from make_dataset import create_hs_ds, load_preproc_dataset\n",
+    "from src.prompts.prompt_loading import load_preproc_dataset\n",
     "\n",
     "# from sklearn.linear_model import LogisticRegression\n",
     "# from sklearn.metrics import f1_score, roc_auc_score, accuracy_score\n",
@@ -73,9 +73,29 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
-   "outputs": [],
+   "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=666)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2023-10-25 13:32:38.035\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-25T13:32:38.035153+0800 INFO changing pad_token_id from 32000 to 0\n",
+      "\u001b[32m2023-10-25 13:32:38.036\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-25T13:32:38.036448+0800 INFO changing padding_side from right to left\n",
+      "\u001b[32m2023-10-25 13:32:38.037\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-25T13:32:38.037328+0800 INFO changing truncation_side from right to left\n"
+     ]
+    }
+   ],
    "source": [
     "# model_name_or_path = \"TheBloke/Wizard-Vicuna-30B-Uncensored-GPTQ\"\n",
     "# model_name_or_path = \"TheBloke/Mistral-7B-Instruct-v0.1-GPTQ\"\n",
@@ -83,7 +103,7 @@
     "\n",
     "batch_size = 2\n",
     "\n",
-    "cfg = ExtractConfig(max_examples=(100, 100), model=model_name_or_path)\n",
+    "cfg = ExtractConfig(max_examples=(100, 100), model=model_name_or_path, max_length=666)\n",
     "print(cfg)\n",
     "\n",
     "model, tokenizer = load_model(model_name_or_path)\n"
@@ -91,7 +111,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -111,12 +131,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
     "# # cache busting for the transformers map and ds steps\n",
-    "# !rm -rf ~/.cache/huggingface/datasets/generator\n"
+    "!rm -rf ~/.cache/huggingface/datasets/generator\n"
    ]
   },
   {
@@ -128,16 +148,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
     "import pickle\n",
-    "N_fit_examples = 10\n",
+    "\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",
+    "def load_rep_reader(model, tokenizer, cfg, N_fit_examples=20, batch_size=2, rep_token = -1, n_difference = 1, direction_method = 'pca'):\n",
     "    \"\"\"\n",
     "    We want one set of interventions per model\n",
     "    \n",
@@ -146,10 +166,7 @@
     "    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",
+    "    if not intervention_f.exists():        \n",
     "        \n",
     "        hidden_layers = list(range(8, model.config.num_hidden_layers, 3))\n",
     "        \n",
@@ -180,15 +197,37 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2023-10-25 13:32:42.237\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mload_rep_reader\u001b[0m:\u001b[36m39\u001b[0m - \u001b[1mLoaded interventions from /media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/data/interventions/TheBloke-WizardCoder-Python-13B-V1.0-GPTQ.pkl\u001b[0m\n",
+      "2023-10-25T13:32:42.237595+0800 INFO Loaded interventions from /media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/data/interventions/TheBloke-WizardCoder-Python-13B-V1.0-GPTQ.pkl\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38]"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "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_fit_examples = 20\n",
+    "N_fit_examples = 10\n",
+    "rep_token = -1\n",
     "\n",
-    "hidden_layers = honesty_rep_reader.directions.keys()\n",
+    "honesty_rep_reader = load_rep_reader(model, tokenizer, cfg, N_fit_examples=N_fit_examples, batch_size=batch_size, rep_token=rep_token)\n",
+    "\n",
+    "hidden_layers = sorted(honesty_rep_reader.directions.keys())\n",
     "hidden_layers\n"
    ]
   },
@@ -201,9 +240,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Generating train split: 0 examples [00:00, ? examples/s]"
+     ]
+    }
+   ],
    "source": [
     "# load dataset\n",
     "ds_name = 'imdb'\n",
@@ -220,6 +267,13 @@
     "dataset_test\n"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -242,6 +296,24 @@
     "# honesty_rep_reader\n"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hidden_layers\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rep_reading_pipeline =  pipeline(\"rep-reading\", model=model, tokenizer=tokenizer)\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -255,9 +327,17 @@
     "    hidden_layers=hidden_layers, \n",
     "    rep_reader=honesty_rep_reader,\n",
     "    batch_size=batch_size, **tokenizer_args)\n",
+    "\n",
     "H_tests[0] # {Batch, layers}\n"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -327,13 +407,6 @@
     "rep_control_pipeline2\n"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -349,7 +422,58 @@
     "    activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer]).to(model.device).half()\n",
     "    \n",
     "\n",
-    "activations_neg = {k:-v for k,v in activations.items()}\n",
+    "activations_neg = {k:-v for k,v in activations.items()}\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# unit test: with multiple input types: single, list, generator, dataset\n",
+    "## single\n",
+    "input_types = {'single':dataset_train[0], 'list':dataset_train[:5], 'generator':iter(dataset_train.select(range(5))), 'dataset':dataset_train.select(range(5)).to_iterable_dataset()}\n",
+    "for name, ds in input_types.items():\n",
+    "    print(f\"==== {name} ====\")\n",
+    "    control_outputs = rep_control_pipeline2(ds, activations=activations, batch_size=2)\n",
+    "    r = list(control_outputs)\n",
+    "    print(len(r))\n",
+    "\n",
+    "\n",
+    "\n",
+    "# control_outputs = rep_control_pipeline2(dataset_train[0], activations=activations)\n",
+    "# list(control_outputs)\n",
+    "\n",
+    "# ## list\n",
+    "# control_outputs = rep_control_pipeline2(dataset_train[:5], activations=activations, batch_size=2)\n",
+    "# list(control_outputs)\n",
+    "\n",
+    "# # generator\n",
+    "# ds = iter(dataset_train.select(range(5)))\n",
+    "# control_outputs = rep_control_pipeline2(ds, activations=activations, batch_size=2)\n",
+    "# list(control_outputs)\n",
+    "\n",
+    "# # dataset\n",
+    "# ds = dataset_train.select(range(5)).to_iterable_dataset()\n",
+    "# control_outputs = rep_control_pipeline2(ds, activations=activations, batch_size=2)\n",
+    "# list(control_outputs)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
     "\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
@@ -366,9 +490,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from torch.utils.data import Dataset\n"
-   ]
+   "source": []
   },
   {
    "cell_type": "code",
@@ -376,8 +498,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# TODO we need to save and cache for many split and datasets\n",
-    "rep_control_pipeline2\n"
+    "from torch.utils.data import Dataset\n"
    ]
   },
   {
@@ -400,13 +521,12 @@
     "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",
+    "def create_hs_ds(ds_name, ds_tokens, pipeline, activations=None, f = None, batch_size=2, split_type=\"train\", debug=False):\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",
+    "    f = str(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",
@@ -422,6 +542,10 @@
     "        batch_size=batch_size,\n",
     "    )\n",
     "    \n",
+    "    if debug:\n",
+    "        # this allow us to debug in a single thread\n",
+    "        pipeline(**gen_kwargs)\n",
+    "    \n",
     "    ds1 = datasets.Dataset.from_generator(\n",
     "        generator=pipeline,\n",
     "        info=datasets.DatasetInfo(\n",
@@ -431,13 +555,15 @@
     "        gen_kwargs=gen_kwargs,\n",
     "        num_proc=1,\n",
     "    )\n",
-    "    return ds1\n",
+    "    logger.info(f\"Created dataset {dataset_name} with {len(ds1)} examples at `{f}`\")\n",
+    "    return ds1, f# TODO we need to save and cache for many split and datasets\n",
+    "rep_control_pipeline2\n",
     "\n",
     "\n",
     "\n",
-    "\n",
-    "ds1 = create_hs_ds('imdb', dataset_train, rep_control_pipeline2, split_type=\"train\")\n",
-    "ds1\n"
+    "ds1, f = create_hs_ds('imdb', dataset_train, rep_control_pipeline2, split_type=\"train\", debug=True)\n",
+    "ds1\n",
+    "\n"
    ]
   },
   {
@@ -445,7 +571,9 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "%debug\n"
+   ]
   },
   {
    "cell_type": "code",

notebooks/make_dataset.py

@@ -48,7 +48,7 @@ from itertools import chain
 import functools
 from src.prompts.prompt_loading import load_prompts
 from src.datasets.scores import scores2choice_probs
-from src.datasets.scores import choice2id, choice2ids
+from src.datasets.scores import choice2id, choice2ids, row_choice_ids
 from src.datasets.intervene import intervention_meta_fn, get_interventions_dict, InterventionDict
 from src.extraction.config import ExtractConfig
 from src.config import root_folder
@@ -223,62 +223,63 @@ def qc_ds(f):
     print(f)
     
     
-    
-def load_preproc_dataset(ds_name: str, cfg: ExtractConfig, tokenizer: PreTrainedTokenizerBase, split_type:str="train", N=None) -> Dataset:
-    """load a preprocessed dataset of tokens."""
-    # TODO refactor out cfg
-    if N is None:
-        N = cfg.max_examples[split_type!="train"]
-    ds_prompts = Dataset.from_generator(
-        load_prompts,
-        gen_kwargs=dict(
-            ds_string=ds_name,
-            num_shots=cfg.num_shots,
-            split_type=split_type,
-            # template_path=template_path,
-            seed=cfg.seed,
-            prompt_format='llama',
-            N=N*3,
-        ),
-    )
+  
+# def load_preproc_dataset(ds_name: str, cfg: ExtractConfig, tokenizer: PreTrainedTokenizerBase, split_type:str="train", N=None) -> Dataset:
+#     """load a preprocessed dataset of tokens."""
+#     # TODO refactor out cfg
+#     if N is None:
+#         N = cfg.max_examples[split_type!="train"]
+#     ds_prompts = Dataset.from_generator(
+#         load_prompts,
+#         gen_kwargs=dict(
+#             ds_string=ds_name,
+#             num_shots=cfg.num_shots,
+#             split_type=split_type,
+#             # template_path=template_path,
+#             seed=cfg.seed,
+#             prompt_format='llama',
+#             N=N*3,
+#         ),
+#     )
 
-    # ## Format prompts
-    # The prompt is the thing we most often have to change and debug. So we do it explicitly here.
-    # We do it as transforms on a huggingface dataset.
-    # In this case we use multishot examples from train, and use the test set to generated the hidden states dataset. We will test generalisation on a whole new dataset.
+#     # ## Format prompts
+#     # The prompt is the thing we most often have to change and debug. So we do it explicitly here.
+#     # We do it as transforms on a huggingface dataset.
+#     # In this case we use multishot examples from train, and use the test set to generated the hidden states dataset. We will test generalisation on a whole new dataset.
 
-    ds_tokens = (
-        ds_prompts
-        .map(
-            lambda ex: tokenizer(
-                ex["question"], padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True,
-                return_tensors="np",
-                return_attention_mask=True,
-                # return_overflowing_tokens=True,
-            ),
-            batched=True,
-            desc='tokenize'
-        )
-        .map(lambda r: {"truncated": np.sum(r["attention_mask"], 0)==cfg.max_length}, desc='truncated')
-        .map(
-            lambda r: {"prompt_truncated": tokenizer.batch_decode(r["input_ids"])},
-            batched=True,
-            desc='prompt_truncated',
-        )
-        .map(lambda r: {'choice_ids': row_choice_ids(r, tokenizer)}, desc='choice_ids')
-    )
+#     ds_tokens = (
+#         ds_prompts
+#         .map(
+#             lambda ex: tokenizer(
+#                 ex["question"], padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True,
+#                 return_tensors="pt",
+#                 return_attention_mask=True,
+#                 # return_overflowing_tokens=True,
+#             ),
+#             batched=True,
+#             desc='tokenize'
+#         )
+#         .map(lambda r: {"truncated": np.sum(r["attention_mask"], 0)==cfg.max_length}, desc='truncated')
+#         .map(
+#             lambda r: {"prompt_truncated": tokenizer.batch_decode(r["input_ids"])},
+#             batched=True,
+#             desc='prompt_truncated',
+#         )
+#         .map(lambda r: {'choice_ids': row_choice_ids(r, tokenizer)}, desc='choice_ids')
+#     )
     
     
+#     ds_tokens = shuffle_dataset_by(ds_tokens, 'example_i')
+#     print('num_rows', ds_tokens.num_rows)
     
-
-    
-    ds_tokens = shuffle_dataset_by(ds_tokens, 'example_i')
-    print('removed truncated rows to leave: num_rows', ds_tokens.num_rows)
-    return ds_tokens
+#     # ## Filter out truncated examples
+#     ds_tokens = ds_tokens.filter(lambda r: not r['prompt_truncated'])
+#     print('num_rows', ds_tokens.num_rows)
+#     return ds_tokens
 
 
-def row_choice_ids(r, tokenizer):
-    return choice2ids([c for c in r['answer_choices']], tokenizer)
+# def row_choice_ids(r, tokenizer):
+#     return choice2ids([c for c in r['answer_choices']], tokenizer)
 
 
 def expand_choices(choices: List[str]) -> Set[str]:

src/datasets/scores.py

@@ -83,3 +83,6 @@ def scores2choice_probs2(logits, choiceids: List[List[int]]):
     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
+
+def row_choice_ids(r, tokenizer):
+    return choice2ids([c for c in r['answer_choices']], tokenizer)

src/prompts/prompt_loading.py

@@ -21,6 +21,11 @@ import functools
 from elk.extraction.balanced_sampler import BalancedSampler, FewShotSampler
 import pandas as pd
 from loguru import logger
+from src.helpers.ds import shuffle_dataset_by
+from src.datasets.scores import choice2id, choice2ids, row_choice_ids
+from src.extraction.config import ExtractConfig
+from src.models.load import verbose_change_param, AutoConfig, AutoTokenizer, AutoModelForCausalLM, PreTrainedTokenizerBase
+
 
 # Local path to the folder containing the templates
 TEMPLATES_FOLDER_PATH = Path(__file__).parent / "templates"
@@ -102,6 +107,7 @@ def load_prompts(
     ds_dict = assert_type(dict, load_dataset(ds_name, config_name or None))
     split_name = select_split(ds_dict, split_type)
 
+    # TODO:, can I make sure it's the same shuffle regardless of length?
     ds = assert_type(Dataset, ds_dict[split_name].shuffle(seed=seed))
     if world_size > 1:
         ds = ds.shard(world_size, rank)
@@ -289,3 +295,58 @@ def _convert_to_prompts(
         raise ValueError(f'Prompt duplicated {dup_count} times! "{maybe_dup}"')
 
     return prompts
+
+
+
+def load_preproc_dataset(ds_name: str, cfg: ExtractConfig, tokenizer: PreTrainedTokenizerBase, split_type:str="train", N=None) -> Dataset:
+    """load a preprocessed dataset of tokens."""
+    # TODO refactor out cfg
+    if N is None:
+        N = cfg.max_examples[split_type!="train"]
+    ds_prompts = Dataset.from_generator(
+        load_prompts,
+        gen_kwargs=dict(
+            ds_string=ds_name,
+            num_shots=cfg.num_shots,
+            split_type=split_type,
+            # template_path=template_path,
+            seed=cfg.seed,
+            prompt_format='llama',
+            N=N*3,
+        ),
+    )
+
+    # ## Format prompts
+    # The prompt is the thing we most often have to change and debug. So we do it explicitly here.
+    # We do it as transforms on a huggingface dataset.
+    # In this case we use multishot examples from train, and use the test set to generated the hidden states dataset. We will test generalisation on a whole new dataset.
+
+    ds_tokens = (
+        ds_prompts
+        .map(
+            lambda ex: tokenizer(
+                ex["question"], padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True,
+                return_tensors="pt",
+                return_attention_mask=True,
+                # return_overflowing_tokens=True,
+            ),
+            batched=True,
+            desc='tokenize'
+        )
+        .map(lambda r: {"truncated": np.sum(r["attention_mask"], 0)==cfg.max_length}, desc='truncated')
+        .map(
+            lambda r: {"prompt_truncated": tokenizer.batch_decode(r["input_ids"])},
+            batched=True,
+            desc='prompt_truncated',
+        )
+        .map(lambda r: {'choice_ids': row_choice_ids(r, tokenizer)}, desc='choice_ids')
+    )
+    
+    
+    ds_tokens = shuffle_dataset_by(ds_tokens, 'example_i')
+    print('num_rows', ds_tokens.num_rows)
+    
+    # ## Filter out truncated examples
+    ds_tokens = ds_tokens.filter(lambda r: not r['truncated'])
+    print('num_rows (after filtering out truncated rows)', ds_tokens.num_rows)
+    return ds_tokens

src/repe/rep_control_pipeline_baukit.py

@@ -10,6 +10,7 @@ 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 einops import rearrange
 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
@@ -53,6 +54,9 @@ def intervention_meta_fn2(
         return intervene(outputs, activations[layer_name])
     else:
         raise ValueError(f"outputs must be tuple or tensor, got {type(outputs)}")
+    
+    
+# def split_outputs(o):
 
 
 class RepControlPipeline2(FeatureExtractionPipeline):
@@ -65,65 +69,99 @@ class RepControlPipeline2(FeatureExtractionPipeline):
         # self.default_class2choiceids = choice2ids(default_class2choices, tokenizer)
 
     def __call__(self, model_inputs, activations=None, **kwargs):
-        if activations is not None:
-            activations_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations.items()})
-            layers_names = [self.layer_name_tmpl.format(i) for i in activations.keys()]
-            edit_fn = partial(intervention_meta_fn2, activations=activations_i)
-            with TraceDict(
-                self.model, layers_names, detach=True, edit_output=edit_fn
-            ) as ret:
+        with torch.no_grad():
+            if activations is not None:
+                activations_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations.items()})
+                layers_names = [self.layer_name_tmpl.format(i) for i in activations.keys()]
+                edit_fn = partial(intervention_meta_fn2, activations=activations_i)
+                with TraceDict(
+                    self.model, layers_names, detach=True, edit_output=edit_fn
+                ) as ret:
+                    outputs = super().__call__(model_inputs, **kwargs)
+            else:
                 outputs = super().__call__(model_inputs, **kwargs)
-        else:
-            outputs = super().__call__(model_inputs, **kwargs)
         return outputs
     
     def preprocess(self, inputs: Dataset, **tokenize_kwargs) -> Dict[str, GenericTensor]:
         # tokenize a batch of inputs
         return_tensors = self.framework
+        
+        # if the pipeline is in "single mode", turn it into a batch
+        if isinstance(inputs['question'], str):
+            inputs = {k: [v] for k, v in inputs.items()}
+            
+        # tokenize if needed
         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
+            inputs = {**inputs, **model_inputs}
+        
+        # to device
+        inputs["input_ids"] = torch.tensor(inputs['input_ids'], dtype=torch.long, device=self.model.device)
+        inputs["attention_mask"] = torch.tensor(inputs['attention_mask'], dtype=torch.bool, device=self.model.device)
+        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}
-        )
+        
+        assert model_inputs['input_ids'].ndim == 2, f"expected input_ids to be (batch, seq), got {model_inputs['input_ids'].shape}"
+        
         self.model.eval()
+        inputs = dict(
+            input_ids=model_inputs['input_ids'],
+            attention_mask=model_inputs['attention_mask'],
+            use_cache=False,
+            output_hidden_states=True,
+            return_dict=True
+        )
         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
+        # hidden states come at as lists of layers, lets concat them
+        model_outputs['hidden_states'] = rearrange(list(model_outputs['hidden_states']), 'l b t h -> b l t h')
+        
+        # batch of outputs and inputs. retain some of the inputs
         model_outputs = {**model_inputs, **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
-        assert isinstance(o, dict) and o['logits'].ndim==3, f"expected dict with logits of shape (batch, seq, vocab), got {o['logits'].shape}"
+        # TODO loop through results and yeild them one at a time
+        for i in range(len(o['input_ids'])):
+            o_i = {k: v[i] for k, v in o.items()}
+            o_i = self.postprocess1(o_i)
+            yield o_i
+        
+    def postprocess1(self, o):
+        assert isinstance(o, dict) and o['logits'].ndim==2, f"expected dict with logits of shape (seq, vocab), got {o['logits'].shape}"
+        # assert o['logits'].shape[0]==1, f"postprocess expected batch size 1, got {o['logits'].shape[0]}"
         # 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()
-        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))
+        
+        # TODO for k in ks: o[k] = o[k].squeeze(0)
+        
+        o["end_logits"] = o["logits"][-1, :].float()
+        
+        input_ids = torch.tensor(o['input_ids'], dtype=torch.long)
+        o["input_truncated"] = self.tokenizer.decode(input_ids)
+        
+        o["truncated"] = torch.tensor(o["attention_mask"]).sum()==self.max_length
+        o["text_ans"] = self.tokenizer.decode(o["end_logits"].argmax(-1))
         
         if 'answer_choices' in o:
             answer_choices = o['answer_choices']
-            if isinstance(answer_choices[0][0], str):
-                answer_choices = [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]
+        o['choice_ids'] = row_choice_ids(answer_choices, self.tokenizer)
         
-        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)
+        p = o['add_ans'] = scores2choice_probs2(o['end_logits'], o['choice_ids']) 
+        o['ans'] = p[1] / (torch.sum(p) + 1e-5)
         
         o = hacky_sanitize_outputs(o)
+        
+        # ah to make a dataset we need to return one at a time, right now it's Dict[str, Batch]. e.g. hiddenstates={layer_1:[2, 555, 5120]....
         return o