diff --git a/mjc_notes.md b/mjc_notes.md
index c6187c8..7ac4d9a 100644
--- a/mjc_notes.md
+++ b/mjc_notes.md
@@ -749,7 +749,7 @@ Where was I?
 
 
 03_ds_TQA... make sure 
-- [ ] true false is rigth, 
+- [ ] true false is right, 
 - [ ] acc
 
 
@@ -760,8 +760,11 @@ TruthfullQA:
 - [ ] OK maybe my own curated set?
 - [ ] Or just make a quick one to test manually...
 
+Refactoring
 - [ ] move common functs to src
-  - [ ] probe
-  - [ ] load model
-  - [ ] get_hidden_states
-  - [ ] get_choices_as_tokens
+  - [x] probe
+  - [x] load model
+  - [x] get_hidden_states
+  - [x] get_choices_as_tokens
+  - [ ] prompt format
+- [ ] get it working :poop:
diff --git a/notebooks/023_mjc_distance_and_direction_loss_96%_conv.ipynb b/notebooks/023_train_prob.ipynb
similarity index 100%
rename from notebooks/023_mjc_distance_and_direction_loss_96%_conv.ipynb
rename to notebooks/023_train_prob.ipynb
diff --git a/notebooks/030_eval.ipynb b/notebooks/030_eval_FIXME.ipynb
similarity index 97%
rename from notebooks/030_eval.ipynb
rename to notebooks/030_eval_FIXME.ipynb
index 70d9c36..9d15b84 100644
--- a/notebooks/030_eval.ipynb
+++ b/notebooks/030_eval_FIXME.ipynb
@@ -242,13 +242,7 @@
     }
    ],
    "source": [
-    "def get_choices_as_tokens(choice_n = \"Negative\", choice_p = \"Positive\"):\n",
-    "    # Note some tokenizer differentiate between \"no\", \"\\nno\", so we sometime need to add whitespace beforehand...\n",
-    "    id_n, id_y = tokenizer(f'\\n{choice_n}', add_special_tokens=True)['input_ids'][-1], tokenizer(f'\\n{choice_p}', add_special_tokens=True)['input_ids'][-1]\n",
-    "    assert tokenizer.decode([id_n])==choice_n\n",
-    "    assert tokenizer.decode([id_y])==choice_p\n",
-    "    # print(tokenizer.decode([id_y]))\n",
-    "    return id_n, id_y"
+    "from src.datasets.hs import get_choices_as_tokens"
    ]
   },
   {
@@ -504,13 +498,6 @@
     "net"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -524,15 +511,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def to_numpy(x):\n",
-    "    if isinstance(x, torch.Tensor):\n",
-    "        # note apache parquet doesn't support half https://github.com/huggingface/datasets/issues/4981\n",
-    "        x = x.detach().cpu().float()\n",
-    "        if x.squeeze().dim()==0:\n",
-    "            return x.item()\n",
-    "        return x.numpy()\n",
-    "    else:\n",
-    "        return x"
+    "from src.helpers.torch import to_numpy"
    ]
   },
   {
diff --git a/notebooks/03_ds.ipynb b/notebooks/03_make_dataset.ipynb
similarity index 90%
rename from notebooks/03_ds.ipynb
rename to notebooks/03_make_dataset.ipynb
index ee5c425..a3d9d09 100644
--- a/notebooks/03_ds.ipynb
+++ b/notebooks/03_make_dataset.ipynb
@@ -213,26 +213,11 @@
     "print(config)\n",
     "config.use_cache = False\n",
     "tokenizer = AutoTokenizer.from_pretrained(model_repo)\n",
-    "model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "None\n"
-     ]
-    }
-   ],
-   "source": [
+    "model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options)\n",
+    "\n",
     "# https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/falcon.py\n",
-    "print(tokenizer.pad_token_id)\n",
     "if tokenizer.pad_token_id is None:\n",
+    "    print(tokenizer.pad_token_id)\n",
     "    tokenizer.pad_token_id = 204 # <unk> https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/alpaca.py\n",
     "tokenizer.padding_side = \"left\""
    ]
@@ -274,23 +259,7 @@
     "N_SHOTS = 3\n",
     "USE_MCDROPOUT = True\n",
     "# dataset_n = 200\n",
-    "N = 8000 # 4000 in 4 hours\n",
-    "\n",
-    "try:\n",
-    "    # num_layers = len(model.model.layers)\n",
-    "    num_layers = model.config.n_layer\n",
-    "    print(num_layers)\n",
-    "except AttributeError:\n",
-    "    try:\n",
-    "        num_layers = len(model.base_model.model.model.layers)\n",
-    "        print(num_layers)\n",
-    "    except:\n",
-    "        num_layers = 10\n",
-    "        \n",
-    "stride = 2\n",
-    "# don't take the first or last layers as they can make it to easy to leak info\n",
-    "extract_layers = tuple(range(2, num_layers-2, stride)) + (num_layers-2,)\n",
-    "extract_layers, num_layers"
+    "N = 8000 # 4000 in 4 hours\n"
    ]
   },
   {
@@ -310,35 +279,10 @@
     }
    ],
    "source": [
-    "# TODO maybe a list of tokens? Maybe the most common from the prompt?\n",
-    "# get the tokens for 0 and 1, we will use these later...\n",
-    "# note that sentancepeice tokenizers have differen't tokens for No and \\nNo.\n",
-    "token_n = \"Negative\"\n",
-    "token_y = \"Positive\"\n",
-    "id_n, id_y = tokenizer(f'\\n{token_n}', add_special_tokens=True)['input_ids'][-1], tokenizer(f'\\n{token_y}', add_special_tokens=True)['input_ids'][-1]\n",
-    "assert tokenizer.decode([id_n])==token_n\n",
-    "assert tokenizer.decode([id_y])==token_y\n",
-    "id_n, id_y"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "'Positive'"
-      ]
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "tokenizer.decode([id_y])"
+    "choices_n = [\"Negative\"]\n",
+    "choices_p = [\"Positive\"]\n",
+    "from src.datasets.hs import get_choices_as_tokens\n",
+    "ids_n, ids_y = get_choices_as_tokens(choices_n, choices_p)"
    ]
   },
   {
@@ -777,25 +721,8 @@
     "    torch.cuda.empty_cache()\n",
     "    gc.collect()\n",
     "    \n",
+    "from src.datasets.dropout import enable_dropout, check_for_dropout\n",
     "\n",
-    "def enable_dropout(model, USE_MCDROPOUT:Union[float,bool]=True):\n",
-    "    \"\"\" Function to enable the dropout layers during test-time \"\"\"\n",
-    "    \n",
-    "    for m in model.modules():\n",
-    "        if m.__class__.__name__.startswith('Dropout'):\n",
-    "            m.train()\n",
-    "            if USE_MCDROPOUT!=True:\n",
-    "                m.p=USE_MCDROPOUT\n",
-    "                # print(m)\n",
-    "                \n",
-    "                \n",
-    "def check_for_dropout(model):\n",
-    "    for m in model.modules():\n",
-    "        if m.__class__.__name__.startswith('Dropout'):\n",
-    "            if m.p>0:\n",
-    "                # print(m)\n",
-    "                return True\n",
-    "    return False\n",
     "    \n",
     "clear_mem()\n",
     "assert check_for_dropout(model), 'model should have dropout modules'\n",
@@ -811,87 +738,79 @@
     "\n",
     "\n",
     "            \n",
-    "def get_hidden_states(model, tokenizer, input_text, layers=extract_layers, truncation_length=999, output_attentions=False, use_mcdropout=USE_MCDROPOUT):\n",
-    "    \"\"\"\n",
-    "    Given a decoder model and some texts, gets the hidden states (in a given layer) on that input texts\n",
-    "    \"\"\"\n",
-    "    if not isinstance(input_text, list):\n",
-    "        input_text = [input_text]\n",
-    "    input_ids = tokenizer(input_text, \n",
-    "                          return_tensors=\"pt\",\n",
-    "                          padding=True,\n",
-    "                            add_special_tokens=True,\n",
-    "                         ).input_ids.to(model.device)\n",
+    "# def get_hidden_states(model, tokenizer, input_text, layers=extract_layers, truncation_length=999, output_attentions=False, use_mcdropout=USE_MCDROPOUT):\n",
+    "#     \"\"\"\n",
+    "#     Given a decoder model and some texts, gets the hidden states (in a given layer) on that input texts\n",
+    "#     \"\"\"\n",
+    "#     if not isinstance(input_text, list):\n",
+    "#         input_text = [input_text]\n",
+    "#     input_ids = tokenizer(input_text, \n",
+    "#                           return_tensors=\"pt\",\n",
+    "#                           padding=True,\n",
+    "#                             add_special_tokens=True,\n",
+    "#                          ).input_ids.to(model.device)\n",
     "    \n",
-    "    # if add_bos_token:\n",
-    "    #     input_ids = input_ids[:, 1:]\n",
+    "#     # if add_bos_token:\n",
+    "#     #     input_ids = input_ids[:, 1:]\n",
     "        \n",
-    "    # Handling truncation: truncate start, not end\n",
-    "    if truncation_length is not None:\n",
-    "        if input_ids.size(1)>truncation_length:\n",
-    "            print('truncating', input_ids.size(1))\n",
-    "        input_ids = input_ids[:, -truncation_length:]\n",
+    "#     # Handling truncation: truncate start, not end\n",
+    "#     if truncation_length is not None:\n",
+    "#         if input_ids.size(1)>truncation_length:\n",
+    "#             print('truncating', input_ids.size(1))\n",
+    "#         input_ids = input_ids[:, -truncation_length:]\n",
     "\n",
-    "    # forward pass\n",
-    "    last_token = -1\n",
-    "    first_token = 0\n",
-    "    with torch.no_grad():\n",
-    "        model.eval()        \n",
-    "        if use_mcdropout: enable_dropout(model, use_mcdropout)\n",
+    "#     # forward pass\n",
+    "#     last_token = -1\n",
+    "#     first_token = 0\n",
+    "#     with torch.no_grad():\n",
+    "#         model.eval()        \n",
+    "#         if use_mcdropout: enable_dropout(model, use_mcdropout)\n",
     "        \n",
-    "        # taken from greedy_decode https://github.com/huggingface/transformers/blob/ba695c1efd55091e394eb59c90fb33ac3f9f0d41/src/transformers/generation/utils.py\n",
-    "        logits_processor = LogitsProcessorList()\n",
-    "        model_kwargs = dict(use_cache=False)\n",
-    "        model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)\n",
-    "        outputs = model.forward(**model_inputs, return_dict=True, output_attentions=output_attentions, output_hidden_states=True)\n",
+    "#         # taken from greedy_decode https://github.com/huggingface/transformers/blob/ba695c1efd55091e394eb59c90fb33ac3f9f0d41/src/transformers/generation/utils.py\n",
+    "#         logits_processor = LogitsProcessorList()\n",
+    "#         model_kwargs = dict(use_cache=False)\n",
+    "#         model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)\n",
+    "#         outputs = model.forward(**model_inputs, return_dict=True, output_attentions=output_attentions, output_hidden_states=True)\n",
     "        \n",
-    "        next_token_logits = outputs.logits[:, last_token, :]\n",
-    "        outputs['scores'] = logits_processor(input_ids, next_token_logits)[:, None,:]\n",
+    "#         next_token_logits = outputs.logits[:, last_token, :]\n",
+    "#         outputs['scores'] = logits_processor(input_ids, next_token_logits)[:, None,:]\n",
     "        \n",
-    "        next_tokens = torch.argmax(outputs['scores'], dim=-1)\n",
-    "        outputs['sequences'] = torch.cat([input_ids, next_tokens], dim=-1)\n",
+    "#         next_tokens = torch.argmax(outputs['scores'], dim=-1)\n",
+    "#         outputs['sequences'] = torch.cat([input_ids, next_tokens], dim=-1)\n",
     "\n",
-    "        # the output is large, so we will just select what we want 1) the first token with[:, 0]\n",
-    "        # 2) selected layers with [layers]\n",
-    "        attentions = None\n",
-    "        if output_attentions:\n",
-    "            # shape is [(batch_size, num_heads, sequence_length, sequence_length)]*num_layers\n",
-    "            # lets take max?\n",
-    "            attentions = [outputs['attentions'][i] for i in layers]\n",
-    "            attentions = [v[:, last_token] for v in attentions]\n",
-    "            attentions = torch.concat(attentions)\n",
+    "#         # the output is large, so we will just select what we want 1) the first token with[:, 0]\n",
+    "#         # 2) selected layers with [layers]\n",
+    "#         attentions = None\n",
+    "#         if output_attentions:\n",
+    "#             # shape is [(batch_size, num_heads, sequence_length, sequence_length)]*num_layers\n",
+    "#             # lets take max?\n",
+    "#             attentions = [outputs['attentions'][i] for i in layers]\n",
+    "#             attentions = [v[:, last_token] for v in attentions]\n",
+    "#             attentions = torch.concat(attentions)\n",
     "        \n",
-    "        hidden_states = torch.stack([outputs['hidden_states'][i] for i in layers], 1)\n",
+    "#         hidden_states = torch.stack([outputs['hidden_states'][i] for i in layers], 1)\n",
     "        \n",
-    "        hidden_states = hidden_states[:, :, last_token] # (batch, layers, past_seq, logits) take just the last token so they are same size\n",
+    "#         hidden_states = hidden_states[:, :, last_token] # (batch, layers, past_seq, logits) take just the last token so they are same size\n",
     "        \n",
-    "        input_truncated = tokenizer.batch_decode(input_ids)\n",
+    "#         input_truncated = tokenizer.batch_decode(input_ids)\n",
     "        \n",
-    "        s = outputs['sequences']\n",
-    "        s = [s[i][len(input_ids[i]):] for i in range(len(s))]\n",
-    "        text_ans = tokenizer.batch_decode(s)\n",
+    "#         s = outputs['sequences']\n",
+    "#         s = [s[i][len(input_ids[i]):] for i in range(len(s))]\n",
+    "#         text_ans = tokenizer.batch_decode(s)\n",
     "\n",
-    "        scores = outputs['scores'][:, first_token].softmax(-1) # for first (and only) token\n",
-    "        prob_n, prob_y = scores[:, [id_n, id_y]].T\n",
-    "        eps = 1e-3\n",
-    "        ans = (prob_y/(prob_n+prob_y+eps))\n",
+    "#         scores = outputs['scores'][:, first_token].softmax(-1) # for first (and only) token\n",
+    "#         prob_n, prob_y = scores[:, [id_n, id_y]].T\n",
+    "#         eps = 1e-3\n",
+    "#         ans = (prob_y/(prob_n+prob_y+eps))\n",
     "    \n",
-    "    out = dict(hidden_states=hidden_states, ans=ans, text_ans=text_ans, input_truncated=input_truncated, input_id_shape=input_ids.shape,\n",
-    "                attentions=attentions, prob_n=prob_n, prob_y=prob_y, scores=outputs['scores'][:, 0], input_text=input_text,\n",
-    "               )\n",
-    "    out = {k:to_numpy(v) for k,v in out.items()}    \n",
-    "    return out\n",
+    "#     out = dict(hidden_states=hidden_states, ans=ans, text_ans=text_ans, input_truncated=input_truncated, input_id_shape=input_ids.shape,\n",
+    "#                 attentions=attentions, prob_n=prob_n, prob_y=prob_y, scores=outputs['scores'][:, 0], input_text=input_text,\n",
+    "#                )\n",
+    "#     out = {k:to_numpy(v) for k,v in out.items()}    \n",
+    "#     return out\n",
     "\n",
     "\n",
-    "def to_numpy(x):\n",
-    "    if isinstance(x, torch.Tensor):\n",
-    "        # note apache parquet doesn't support half https://github.com/huggingface/datasets/issues/4981\n",
-    "        x = x.detach().cpu().float()\n",
-    "        if x.squeeze().dim()==0:\n",
-    "            return x.item()\n",
-    "        return x.numpy()\n",
-    "    else:\n",
-    "        return x"
+    "from src.helpers.torch import to_numpy"
    ]
   },
   {
@@ -908,8 +827,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def md5hash(s: bytes) -> str:\n",
-    "    return hashlib.md5(s).hexdigest()"
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from src.datasets.hs import ExtractHiddenStates\n",
+    "from src.datasets.batch import batch_hidden_states\n",
+    "ehs = ExtractHiddenStates(model, tokenizer, choices_n, choices_y)\n",
+    "batch_hidden_states(ehs,)\n"
    ]
   },
   {
@@ -919,69 +849,71 @@
    "outputs": [],
    "source": [
     "\n",
-    "def batch_hidden_states(prompt_fn=format_imdbs_multishot, model=model, tokenizer=tokenizer, data=data, n=100, batch_size=2, version_options=['lie', 'truth'], mcdropout=True):\n",
-    "    \"\"\"\n",
-    "    Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.\n",
-    "    Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)\n",
-    "    with the ground truth labels\n",
+    "# def batch_hidden_states(prompt_fn=format_imdbs_multishot, model=model, tokenizer=tokenizer, data=data, n=100, batch_size=2, version_options=['lie', 'truth'], mcdropout=True):\n",
+    "#     \"\"\"\n",
+    "#     Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.\n",
+    "#     Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)\n",
+    "#     with the ground truth labels\n",
     "    \n",
-    "    This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency\n",
-    "    \"\"\"\n",
-    "    # setup\n",
-    "    model.eval()\n",
+    "#     This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency\n",
+    "#     \"\"\"\n",
     "    \n",
-    "    ds_subset = data.shuffle(seed=42).select(range(n))\n",
-    "    dl = DataLoader(ds_subset, batch_size=batch_size, shuffle=True)\n",
-    "    for i, batch in enumerate(tqdm(dl, desc='get hidden states')):\n",
-    "        titles, contents, true_labels =  batch[\"title\"], batch[\"content\"], batch[\"label\"]\n",
-    "        texts = [format_review(t, c) for t,c in zip(titles, contents)]\n",
-    "        nn = len(texts)\n",
-    "        index = i*batch_size+np.arange(nn)\n",
-    "        for version in version_options:\n",
-    "            versions = [version]*nn\n",
-    "            q, info = prompt_fn(texts, answers=true_labels, versions=versions)\n",
-    "            if i==0:\n",
-    "                assert len(texts)==len(prompt_fn(texts)[0]), 'make sure the prompt function can handle a list of text'\n",
+    "    \n",
+    "#     # setup\n",
+    "#     model.eval()\n",
+    "    \n",
+    "#     ds_subset = data.shuffle(seed=42).select(range(n))\n",
+    "#     dl = DataLoader(ds_subset, batch_size=batch_size, shuffle=True)\n",
+    "#     for i, batch in enumerate(tqdm(dl, desc='get hidden states')):\n",
+    "#         titles, contents, true_labels =  batch[\"title\"], batch[\"content\"], batch[\"label\"]\n",
+    "#         texts = [format_review(t, c) for t,c in zip(titles, contents)]\n",
+    "#         nn = len(texts)\n",
+    "#         index = i*batch_size+np.arange(nn)\n",
+    "#         for version in version_options:\n",
+    "#             versions = [version]*nn\n",
+    "#             q, info = prompt_fn(texts, answers=true_labels, versions=versions)\n",
+    "#             if i==0:\n",
+    "#                 assert len(texts)==len(prompt_fn(texts)[0]), 'make sure the prompt function can handle a list of text'\n",
     "            \n",
-    "            # different due to dropout\n",
-    "            # set_seeds(i*10)\n",
-    "            hs1 = get_hidden_states(model, tokenizer, q, use_mcdropout=mcdropout)\n",
-    "            # set_seeds(i*10+1)\n",
-    "            if mcdropout:\n",
-    "                hs2 = get_hidden_states(model, tokenizer, q, use_mcdropout=mcdropout)\n",
+    "#             # different due to dropout\n",
+    "#             # set_seeds(i*10)\n",
+    "#             hs1 = ehs.get_hidden_states(q, use_mcdropout=mcdropout)\n",
+    "#             # set_seeds(i*10+1)\n",
+    "#             if mcdropout:\n",
+    "#                 hs2 = ehs.get_hidden_states(q, use_mcdropout=mcdropout)\n",
     "                \n",
-    "                # QC\n",
-    "                if i==0:\n",
-    "                    eps=1e-5\n",
-    "                    mpe = lambda x,y: np.mean(np.abs(x-y)/(np.abs(x)+np.abs(y)+eps))\n",
-    "                    a,b=hs2['hidden_states'],hs1['hidden_states']\n",
-    "                    assert mpe(a,b)>eps, \"the hidden state pairs should be different but are not. Check model.config.use_cache==False, check this model has dropout in it's arch\"\n",
+    "#                 # QC\n",
+    "#                 if i==0:\n",
+    "#                     eps=1e-5\n",
+    "#                     mpe = lambda x,y: np.mean(np.abs(x-y)/(np.abs(x)+np.abs(y)+eps))\n",
+    "#                     a,b=hs2['hidden_states'],hs1['hidden_states']\n",
+    "#                     assert mpe(a,b)>eps, \"the hidden state pairs should be different but are not. Check model.config.use_cache==False, check this model has dropout in it's arch\"\n",
     "                    \n",
-    "                    assert ((hs1['prob_y']+hs1['prob_n'])>0.5).all(), \"your chosen binary answers should take up a lot of the prob space, otherwise choose differen't tokens\"\n",
-    "            else:\n",
-    "                hs2 = hs1\n",
+    "#                     assert ((hs1['prob_y']+hs1['prob_n'])>0.5).all(), \"your chosen binary answers should take up a lot of the prob space, otherwise choose differen't tokens\"\n",
+    "#             else:\n",
+    "#                 hs2 = hs1\n",
     "\n",
     "\n",
-    "            for j in range(nn):\n",
-    "                yield dict(\n",
-    "                    hs1=hs1['hidden_states'][j],\n",
-    "                    ans1=hs1[\"ans\"][j],\n",
+    "#             for j in range(nn):\n",
+    "#                 yield dict(\n",
+    "#                     hs1=hs1['hidden_states'][j],\n",
+    "#                     ans1=hs1[\"ans\"][j],\n",
     "                    \n",
-    "                    hs2=hs2['hidden_states'][j],\n",
-    "                    ans2=hs2[\"ans\"][j],                    \n",
+    "#                     hs2=hs2['hidden_states'][j],\n",
+    "#                     ans2=hs2[\"ans\"][j],                    \n",
     "                    \n",
-    "                    true=true_labels[j].item(),\n",
-    "                    index=index[j],\n",
-    "                    version=version,\n",
-    "                    info=info[j],\n",
+    "#                     true=true_labels[j].item(),\n",
+    "#                     index=index[j],\n",
+    "#                     version=version,\n",
+    "#                     info=info[j],\n",
     "                    \n",
-    "                    # optional/debug\n",
-    "                    input_truncated=hs1['input_truncated'][j], # the question after truncating\n",
-    "                    prob_y=hs1['prob_y'][j],\n",
-    "                    prob_n=hs1['prob_n'][j],\n",
-    "                    text_ans = hs1['text_ans'][j],\n",
-    "                    input_text=hs1['input_text'][j],\n",
-    "                )"
+    "#                     # optional/debug\n",
+    "#                     input_truncated=hs1['input_truncated'][j], # the question after truncating\n",
+    "#                     prob_y=hs1['prob_y'][j],\n",
+    "#                     prob_n=hs1['prob_n'][j],\n",
+    "#                     text_ans = hs1['text_ans'][j],\n",
+    "#                     input_text=hs1['input_text'][j],\n",
+    "#                 )"
    ]
   },
   {
@@ -1378,6 +1310,9 @@
     }
    ],
    "source": [
+    "def md5hash(s: bytes) -> str:\n",
+    "    return hashlib.md5(s).hexdigest()\n",
+    "\n",
     "# unique hash\n",
     "def get_unique_config_name(prompt_fn, model, tokenizer, data, N):\n",
     "    \"\"\"\n",
diff --git a/src/datasets/batch.py b/src/datasets/batch.py
new file mode 100644
index 0000000..b4c6a23
--- /dev/null
+++ b/src/datasets/batch.py
@@ -0,0 +1,67 @@
+
+from tqdm.auto import tqdm
+from src.datasets.hs import ExtractHiddenStates
+from torch.utils.data import DataLoader
+import numpy as np
+
+def batch_hidden_states(ehs: ExtractHiddenStates, prompt_fn=format_imdbs_multishot, data=data, n=100, batch_size=2, version_options=['lie', 'truth'], mcdropout=True):
+    """
+    Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.
+    Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)
+    with the ground truth labels
+    
+    This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency
+    """
+    
+    ds_subset = data.shuffle(seed=42).select(range(n))
+    dl = DataLoader(ds_subset, batch_size=batch_size, shuffle=True)
+    for i, batch in enumerate(tqdm(dl, desc='get hidden states')):
+        titles, contents, true_labels =  batch["title"], batch["content"], batch["label"]
+        texts = [format_review(t, c) for t,c in zip(titles, contents)]
+        nn = len(texts)
+        index = i*batch_size+np.arange(nn)
+        for version in version_options:
+            versions = [version]*nn
+            q, info = prompt_fn(texts, answers=true_labels, versions=versions)
+            if i==0:
+                assert len(texts)==len(prompt_fn(texts)[0]), 'make sure the prompt function can handle a list of text'
+            
+            # different due to dropout
+            # set_seeds(i*10)
+            hs1 = ehs.get_hidden_states(q, use_mcdropout=mcdropout)
+            # set_seeds(i*10+1)
+            if mcdropout:
+                hs2 = ehs.get_hidden_states(q, use_mcdropout=mcdropout)
+                
+                # QC
+                if i==0:
+                    eps=1e-5
+                    mpe = lambda x,y: np.mean(np.abs(x-y)/(np.abs(x)+np.abs(y)+eps))
+                    a,b=hs2['hidden_states'],hs1['hidden_states']
+                    assert mpe(a,b)>eps, "the hidden state pairs should be different but are not. Check model.config.use_cache==False, check this model has dropout in it's arch"
+                    
+                    assert ((hs1['prob_y']+hs1['prob_n'])>0.5).all(), "your chosen binary answers should take up a lot of the prob space, otherwise choose differen't tokens"
+            else:
+                hs2 = hs1
+
+
+            for j in range(nn):
+                yield dict(
+                    hs1=hs1['hidden_states'][j],
+                    ans1=hs1["ans"][j],
+                    
+                    hs2=hs2['hidden_states'][j],
+                    ans2=hs2["ans"][j],                    
+                    
+                    true=true_labels[j].item(),
+                    index=index[j],
+                    version=version,
+                    info=info[j],
+                    
+                    # optional/debug
+                    input_truncated=hs1['input_truncated'][j], # the question after truncating
+                    prob_y=hs1['prob_y'][j],
+                    prob_n=hs1['prob_n'][j],
+                    text_ans = hs1['text_ans'][j],
+                    input_text=hs1['input_text'][j],
+                )
diff --git a/src/datasets/dm.py b/src/datasets/dm.py
index 02a0895..a0f8e5a 100644
--- a/src/datasets/dm.py
+++ b/src/datasets/dm.py
@@ -2,12 +2,12 @@ import torch
 import torch.nn as nn
 import lightning as pl
 import pandas as pd
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset, DataLoader, TensorDataset
 from src.datasets.load import ds2df
 
 def make_y(df):
     # label: is ans2 more true than ans1
-    # so we ask does ans2 have greater probabiliy on "positive" than ans1
+    # so we ask does ans2 have greater probability on "positive" than ans1
     # then, when the right answer is negative we swap the sign
     true_switch_sign = df.true_answer*2-1
     distance = (df.ans2-df.ans1) * true_switch_sign
diff --git a/src/datasets/dropout.py b/src/datasets/dropout.py
new file mode 100644
index 0000000..cf229d7
--- /dev/null
+++ b/src/datasets/dropout.py
@@ -0,0 +1,20 @@
+
+from typing import Union
+
+def enable_dropout(model, USE_MCDROPOUT:Union[float,bool]=True):
+    """ Function to enable the dropout layers during test-time """
+    
+    for m in model.modules():
+        if m.__class__.__name__.startswith('Dropout'):
+            m.train()
+            if USE_MCDROPOUT!=True:
+                m.p=USE_MCDROPOUT
+                
+                
+def check_for_dropout(model, verbose=False):
+    for m in model.modules():
+        if m.__class__.__name__.startswith('Dropout'):
+            if m.p>0:
+                if verbose: print(m)
+                return True
+    return False
diff --git a/src/datasets/hs.py b/src/datasets/hs.py
new file mode 100644
index 0000000..f3c876a
--- /dev/null
+++ b/src/datasets/hs.py
@@ -0,0 +1,177 @@
+from dataclasses import dataclass
+import lightning as pl
+import torch
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    AutoModelForMaskedLM,
+    AutoModelForCausalLM,
+    AutoConfig,
+    AutoModel,
+    PreTrainedTokenizer,
+    PreTrainedModel
+)
+from typing import Optional, List, Tuple
+from transformers import LogitsProcessorList
+
+from src.helpers.torch import to_numpy
+from src.datasets.dropout import enable_dropout
+
+
+def get_choices_as_tokens(
+    tokenizer, choice_n: List[str] = ["Negative"], choice_p: List[str] = ["Positive"]
+) -> Tuple[List[int], List[int]]:
+    # Note some tokenizer differentiate between "no", "\nno", so we sometime need to add whitespace beforehand...
+    ids_n = []
+    for c in choice_n:
+        id_ = tokenizer(f"\n{c}", add_special_tokens=True)["input_ids"][-1]
+        ids_n.append(id_)
+        assert tokenizer.decode([id_]) == c
+
+    ids_y = []
+    for c in choice_n:
+        id_ = tokenizer(f"\n{c}", add_special_tokens=True)["input_ids"][-1]
+        ids_y.append(id_)
+        assert tokenizer.decode([id_]) == c
+
+    return ids_n, ids_y
+
+
+@dataclass
+class ExtractHiddenStates:
+    model: PreTrainedModel
+    tokenizer: PreTrainedTokenizer
+    layer_stride: int = 1
+    layer_padding: int = 2
+    truncation_length = 999
+    choices_n: List[str] = ["No"]
+    choices_p: List[str] = ["Yes"]
+
+    def start(self):
+        self.ids_n, self.ids_y = get_choices_as_tokens(
+            self.tokenizer, self.choices_n, self.choices_p
+        )
+
+    def get_hidden_states(
+        self,
+        input_text,
+        truncation_length=999,
+        output_attentions=False,
+        use_mcdropout=True,
+    ):
+        """
+        Given a decoder model and some texts, gets the hidden states (in a given layer) on that input texts
+        """
+        if not isinstance(input_text, list):
+            input_text = [input_text]
+        input_ids = self.tokenizer(
+            input_text,
+            return_tensors="pt",
+            padding=True,
+            add_special_tokens=True,
+        ).input_ids.to(self.model.device)
+
+        # Handling truncation: truncate start, not end
+        if truncation_length is not None:
+            if input_ids.size(1) > truncation_length:
+                print("truncating", input_ids.size(1))
+            input_ids = input_ids[:, -truncation_length:]
+
+        # forward pass
+        last_token = -1
+        first_token = 0
+        with torch.no_grad():
+            self.model.eval()
+            if use_mcdropout:
+                enable_dropout(self.model, use_mcdropout)
+
+            # taken from greedy_decode https://github.com/huggingface/transformers/blob/ba695c1efd55091e394eb59c90fb33ac3f9f0d41/src/transformers/generation/utils.py
+            logits_processor = LogitsProcessorList()
+            model_kwargs = dict(use_cache=False)
+            model_inputs = self.model.prepare_inputs_for_generation(
+                input_ids, **model_kwargs
+            )
+            outputs = self.model.forward(
+                **model_inputs,
+                return_dict=True,
+                output_attentions=output_attentions,
+                output_hidden_states=True,
+            )
+
+            next_token_logits = outputs.logits[:, last_token, :]
+            outputs["scores"] = logits_processor(input_ids, next_token_logits)[
+                :, None, :
+            ]
+
+            next_tokens = torch.argmax(outputs["scores"], dim=-1)
+            outputs["sequences"] = torch.cat([input_ids, next_tokens], dim=-1)
+
+            # the output is large, so we will just select what we want 1) the first token with[:, 0]
+            # 2) selected layers with [layers]
+            layers = self.get_layer_selection(outputs)
+            attentions = None
+            layers = range(
+                self.layer_padding,
+                len(outputs["attentions"]) - self.layer_padding,
+                self.layer_stride,
+            )
+            if output_attentions:
+                # shape is [(batch_size, num_heads, sequence_length, sequence_length)]*num_layers
+                # lets take max?
+                attentions = [outputs["attentions"][i] for i in layers]
+                attentions = [v[:, last_token] for v in attentions]
+                attentions = torch.concat(attentions)
+
+            hidden_states = torch.stack(
+                [outputs["hidden_states"][i] for i in layers], 1
+            )
+
+            hidden_states = hidden_states[
+                :, :, last_token
+            ]  # (batch, layers, past_seq, logits) take just the last token so they are same size
+
+            input_truncated = self.tokenizer.batch_decode(input_ids)
+
+            s = outputs["sequences"]
+            s = [s[i][len(input_ids[i]) :] for i in range(len(s))]
+            text_ans = self.tokenizer.batch_decode(s)
+
+            scores = outputs["scores"][:, first_token].softmax(
+                -1
+            )  # for first (and only) token
+            # prob_n, prob_y = scores[:, [id_n, id_y]].T
+            prob_n = scores[:, self.ids_n]
+            prob_y = scores[:, self.ids_y]
+            eps = 1e-3
+            ans = (prob_y / (prob_n + prob_y + eps)).sum(1)
+
+        out = dict(
+            hidden_states=hidden_states,
+            ans=ans,
+            text_ans=text_ans,
+            input_truncated=input_truncated,
+            input_id_shape=input_ids.shape,
+            attentions=attentions,
+            prob_n=prob_n,
+            prob_y=prob_y,
+            scores=outputs["scores"][:, 0],
+            input_text=input_text,
+        )
+        out = {k: to_numpy(v) for k, v in out.items()}
+        return out
+
+    def get_layer_selection(self, outputs):
+        """Sometimes we don't want to save all layers.
+
+        Typically we can skip some to save space (stride). We might also want to ignore the first and last ones (padding) to avoid data leakage.
+
+        See https://www.lesswrong.com/posts/bWxNPMy5MhPnQTzKz/what-discovering-latent-knowledge-did-and-did-not-find-4
+        """
+        return range(
+            self.layer_padding,
+            len(outputs["attentions"]) - self.layer_padding,
+            self.layer_stride,
+        )
+
+
+
diff --git a/src/datasets/load.py b/src/datasets/load.py
index 2c10cc2..abcb77c 100644
--- a/src/datasets/load.py
+++ b/src/datasets/load.py
@@ -1,3 +1,6 @@
+import numpy as np
+import pandas as pd
+
 def rows_item(row):
     """
     transform a row by turning singe dim arrays into items
diff --git a/src/helpers/torch.py b/src/helpers/torch.py
new file mode 100644
index 0000000..2522603
--- /dev/null
+++ b/src/helpers/torch.py
@@ -0,0 +1,14 @@
+import torch
+
+def to_numpy(x):
+    """
+    Trys to convert torch to numpy and if possible a single item
+    """
+    if isinstance(x, torch.Tensor):
+        # note apache parquet doesn't support half https://github.com/huggingface/datasets/issues/4981
+        x = x.detach().cpu().float()
+        if x.squeeze().dim()==0:
+            return x.item()
+        return x.numpy()
+    else:
+        return x