diff --git a/mjc_notes.md b/mjc_notes.md
index c9f1aeb..1d42c81 100644
--- a/mjc_notes.md
+++ b/mjc_notes.md
@@ -69,3 +69,19 @@ I guess this shows they they trained whether the text it read is true... because
 
 bug: so there are two no tokens... wtf
 the model only uses one! wtf!
+
+oh it's _No vs No. One is the start of a new word. E.g. " No" and "\nNo" are _No. But "No" is "No"
+
+
+Q for elk
+- why int16 quant of hidden states? oh for the datasets package! I see https://github.com/EleutherAI/elk/issues/208
+- why use forward not generate?
+  - forward is what it thought of the last token?
+  - generate is what it thinks of the generated tokens, conditional on what it read. 
+- The later seems much closer to mind reading, and much closer to what we want
+- additionally I've made this change in a PR and get X results
+
+# 2023-05-29 07:09:38
+
+Note we are using normalized for sklearn and it seems easy?
+but unnorm for CCS, hmm
diff --git a/notebooks/004_mjc_CCS_v2.ipynb b/notebooks/004_mjc_CCS_v2.ipynb
index d6612da..88b9d1b 100644
--- a/notebooks/004_mjc_CCS_v2.ipynb
+++ b/notebooks/004_mjc_CCS_v2.ipynb
@@ -61,78 +61,41 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n",
-      "===================================BUG REPORT===================================\n",
-      "Welcome to bitsandbytes. For bug reports, please run\n",
-      "\n",
-      "python -m bitsandbytes\n",
-      "\n",
-      " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
-      "================================================================================\n",
-      "bin /home/ubuntu/mambaforge/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda117.so\n",
-      "CUDA SETUP: CUDA runtime path found: /home/ubuntu/mambaforge/envs/dlk2/lib/libcudart.so.11.0\n",
-      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
-      "CUDA SETUP: Detected CUDA version 117\n",
-      "CUDA SETUP: Loading binary /home/ubuntu/mambaforge/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda117.so...\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/ubuntu/mambaforge/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: Found duplicate ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] files: {PosixPath('/home/ubuntu/mambaforge/envs/dlk2/lib/libcudart.so.11.0'), PosixPath('/home/ubuntu/mambaforge/envs/dlk2/lib/libcudart.so')}.. We'll flip a coin and try one of these, in order to fail forward.\n",
-      "Either way, this might cause trouble in the future:\n",
-      "If you get `CUDA error: invalid device function` errors, the above might be the cause and the solution is to make sure only one ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] in the paths that we search based on your env.\n",
-      "  warn(msg)\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f444c7277f164769a7c1893f2e14db58",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Loading checkpoint shards:   0%|          | 0/7 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
+    "\n",
     "\n",
     "# 7B\n",
     "# model_repo = \"Neko-Institute-of-Science/LLaMA-7B-HF\"\n",
     "# lora_repo = \"chansung/gpt4-alpaca-lora-7b\"\n",
     "\n",
-    "# 13B\n",
+    "# 13B these work with a batch size of 14 and 2-shot\n",
     "model_repo = \"Neko-Institute-of-Science/LLaMA-13B-HF\"\n",
     "lora_repo = \"chansung/gpt4-alpaca-lora-13b\"\n",
     "\n",
     "model_repo = \"TheBloke/Wizard-Vicuna-13B-Uncensored-HF\"\n",
     "lora_repo = None\n",
     "\n",
-    "# 30B\n",
-    "model_repo = \"TheBloke/OpenAssistant-SFT-7-Llama-30B-HF\"\n",
+    "# 30B - these work but with batch size <=2 & 2-shot\n",
+    "# model_repo = \"TheBloke/OpenAssistant-SFT-7-Llama-30B-HF\"\n",
     "# model_repo = \"ausboss/llama-30b-supercot\"\n",
     "# model_repo= \"timdettmers/guanaco-33b-merged\"\n",
-    "lora_repo = None\n",
+    "# lora_repo = None\n",
     "\n",
     "\n",
     "model_options = dict(\n",
     "    device_map=\"auto\", \n",
     "    load_in_4bit=True,\n",
     "    torch_dtype=torch.float16,\n",
-    "    # bnb_4bit_use_double_quant=True,\n",
     ")\n",
     "\n",
     "    \n",
@@ -165,7 +128,15 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "batch_size = 12\n",
+    "dataset_n = 2000\n",
+    "\n",
+    "num_layers = len(model.model.layers)\n",
+    "stride = 4\n",
+    "extract_layers = (0,) + tuple(range(1, num_layers + 1, stride))\n",
+    "extract_layers"
+   ]
   },
   {
    "cell_type": "code",
@@ -233,9 +204,9 @@
    "source": [
     "def format_imdb_multishot(input, question, response=\"\", lie=False):\n",
     "    a = format_single_shot(\"Although Hypnotic isn't without glimmers of inspiration, the ultimate effect of this often clunky crime caper will be to leave you feeling rather sleepy.\", 0^lie, \"Yes\", include_prefix=True, lie=lie)\n",
-    "    b = format_single_shot(\"Although Bewitched isn't without it's downsides, well to tell the truth, it's all downsides.\", 1^lie, \"No\")\n",
+    "    # b = format_single_shot(\"Although Bewitched isn't without it's downsides, well to tell the truth, it's all downsides.\", 1^lie, \"No\")\n",
     "    c = format_single_shot(input, question, response)\n",
-    "    return \"\\n\\n\".join([a, b, c])"
+    "    return \"\\n\\n\".join([a, c])"
    ]
   },
   {
@@ -324,6 +295,27 @@
     "    return wrap\n"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# from transformers import GenerationConfig\n",
+    "# # from https://github.com/deep-diver/LLM-As-Chatbot/blob/main/configs/response_configs/default.yaml\n",
+    "# # https://github.com/oobabooga/text-generation-webui/blob/main/presets/LLaMA-Precise.txt\n",
+    "# generation_config = GenerationConfig(\n",
+    "#     temperature=0.7,\n",
+    "#     top_p=0.1,\n",
+    "#     top_k=40,\n",
+    "#     num_beams=1,\n",
+    "#     use_cache=True,\n",
+    "#     repetition_penalty=1.18,\n",
+    "#     max_new_tokens=2,\n",
+    "#     do_sample=False,\n",
+    "# )"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -344,7 +336,7 @@
     ")\n",
     "\n",
     "\n",
-    "def get_hidden_states(model, tokenizer, input_text, layers=[2, -2], add_bos_token=1, truncation_length=400, output_attentions=False):\n",
+    "def get_hidden_states(model, tokenizer, input_text, layers=extract_layers, add_bos_token=1, truncation_length=400, output_attentions=False):\n",
     "    \"\"\"\n",
     "    Given a decoder model and some texts, gets the hidden states (in a given layer) on that input texts\n",
     "    \"\"\"\n",
@@ -412,7 +404,7 @@
    "outputs": [],
    "source": [
     "@cache_strargs_kwargs\n",
-    "def batch_hidden_states(model, tokenizer, data, prompt_fn, n=100, layers=[2, -2], batch_size=12):\n",
+    "def batch_hidden_states(model, tokenizer, data, prompt_fn, n=100, layers=extract_layers, batch_size=12):\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",
@@ -447,80 +439,6 @@
     "    return res"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from transformers import GenerationConfig\n",
-    "# from https://github.com/deep-diver/LLM-As-Chatbot/blob/main/configs/response_configs/default.yaml\n",
-    "# https://github.com/oobabooga/text-generation-webui/blob/main/presets/LLaMA-Precise.txt\n",
-    "generation_config = GenerationConfig(\n",
-    "    temperature=0.7,\n",
-    "    top_p=0.1,\n",
-    "    top_k=40,\n",
-    "    num_beams=1,\n",
-    "    use_cache=True,\n",
-    "    repetition_penalty=1.18,\n",
-    "    max_new_tokens=2,\n",
-    "    do_sample=False,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## QC: Try a single pass"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "test_text_pairs = [\n",
-    "    # text, sentiment\n",
-    "    ['This movie was trash burger. It was a very bad movie.', 0],\n",
-    "    [\"This movie changed my life, I've watched it over 5 times and shown it to my entire family\", 1],\n",
-    "    [\"\"\"Lifetime did it again. Can we say stupid? I couldn't wait for it to end. The plot was senseless. The acting was terrible! Especially by the teenagers. The story has been played a thousand times! Are we just desperate to give actors a job? The previews were attractive and I was really looking for a good thriller.Once in awhile lifetime comes up with a good movie, this isn't one of them. Unless one has nothing else to do I would avoid this one at all cost. This was a waste of two hours of my life. Can I get them back? I would have rather scraped my face against a brick wall for two hours then soaked it in peroxide. That would have been more entertaining.\"\"\", 0],\n",
-    "    [\"I can't remember many films where a bumbling idiot of a hero was so funny throughout. Leslie Cheung is such the antithesis of a hero that he's too dense to be seduced by a gorgeous vampire... I had the good luck to see it on a big screen, and to find a video to watch again and again. 9/10\", 1],\n",
-    "    [\"The little girl Desi is so adorable... I cant think of a more beautiful story then this one here. It will make you cry, laugh, and believe. Knowing that this was based on a true story just made me gasp and it also made me realize that there are nice people out there. Great cast and an overall great movie.\", 1],    \n",
-    "]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# attempt at meta example....\n",
-    "\n",
-    "## Params\n",
-    "lie=0\n",
-    "question=0\n",
-    "\n",
-    "i = 3\n",
-    "text = [test_text_pairs[i][0]]\n",
-    "answer = test_text_pairs[i][1]\n",
-    "\n",
-    "## run\n",
-    "neg = get_hidden_states(model, tokenizer, format_imdbs_multishot(text, question, lie=lie))\n",
-    "\n",
-    "## display\n",
-    "print(neg['text_q'][0])\n",
-    "print('='*80)\n",
-    "desired_ans=(question==answer)^lie\n",
-    "print(f\"question=q={question}, answer=a={answer}, lie=l={lie}. (q*a)^l==(({question}*{answer})^{lie}=={desired_ans}) \")\n",
-    "print(f'[public textual answer should be `{\"Yes\" if (question==answer)^lie  else \"No\"}` for this to be a {\"lie\" if lie else \"truth\"}:]')\n",
-    "print(neg['text_ans'][0])\n",
-    "print(f'[public numeric answer should be {\">50%\" if (desired_ans) else \"<50%\"}')\n",
-    "print(f\"{neg['ans'][0]:2.2%}\")"
-   ]
-  },
   {
    "attachments": {},
    "cell_type": "markdown",
@@ -529,6 +447,13 @@
     "## Lightning DataModule"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -544,6 +469,7 @@
     "                 dataset_name=\"amazon_polarity\",\n",
     "                 batch_size=2,\n",
     "                 n=6000,\n",
+    "                 layers=extract_layers,\n",
     "                ):\n",
     "        super().__init__()\n",
     "        self.model = model\n",
@@ -564,7 +490,7 @@
     "\n",
     "        # in ELK they cache as a huggingface dataset\n",
     "        self.neg_hs, self.pos_hs, self.y, self.all_neg_ans, self.all_pos_ans = batch_hidden_states(\n",
-    "            self.model, self.tokenizer, self.dataset, self.prompt_fn, n=h.n, layers=[2, -2], batch_size=h.batch_size)\n",
+    "            self.model, self.tokenizer, self.dataset, self.prompt_fn, n=h.n, layers=h.layers, batch_size=h.batch_size)\n",
     "\n",
     "        # let's create a simple 50/50 train split (the data is already randomized)\n",
     "        n = len(self.y)\n",
@@ -578,19 +504,7 @@
     "            val_split:test_split], self.y[val_split:test_split]\n",
     "        neg_hs_test, pos_hs_test, y_test = self.neg_hs[test_split:],self. pos_hs[\n",
     "            test_split:], self.y[test_split:]\n",
-    "\n",
-    "        # for simplicity we can just take the difference between positive and negative hidden states\n",
-    "        # (concatenating also works fine)\n",
-    "        self.x_train = neg_hs_train - pos_hs_train\n",
-    "        self.x_val = neg_hs_val - pos_hs_val\n",
-    "        self.x_test = neg_hs_test - pos_hs_test\n",
-    "\n",
-    "        # normalize\n",
-    "        self.scaler = RobustScaler()\n",
-    "        self.scaler.fit(self.x_train)\n",
-    "        self.x_train = self.scaler.transform(self.x_train)\n",
-    "        self.x_val = self.scaler.transform(self.x_val)\n",
-    "        self.x_test = self.scaler.transform(self.x_test)\n",
+    "        \n",
     "\n",
     "        self.ds_train = TensorDataset(torch.from_numpy(neg_hs_train).float(),\n",
     "                                      torch.from_numpy(pos_hs_train).float(),\n",
@@ -604,6 +518,19 @@
     "                                     torch.from_numpy(pos_hs_test).float(),\n",
     "                                     torch.from_numpy(y_test).float())\n",
     "\n",
+    "        # for simplicity and sklearn we can just take the difference between positive and negative hidden states\n",
+    "        # (concatenating also works fine)\n",
+    "        self.x_train = neg_hs_train - pos_hs_train\n",
+    "        self.x_val = neg_hs_val - pos_hs_val\n",
+    "        self.x_test = neg_hs_test - pos_hs_test\n",
+    "\n",
+    "        # normalize\n",
+    "        self.scaler = RobustScaler()\n",
+    "        self.scaler.fit(self.x_train)\n",
+    "        self.x_train = self.scaler.transform(self.x_train)\n",
+    "        self.x_val = self.scaler.transform(self.x_val)\n",
+    "        self.x_test = self.scaler.transform(self.x_test)\n",
+    "\n",
     "    def train_dataloader(self):\n",
     "        return DataLoader(self.ds_train,\n",
     "                          batch_size=self.hparams.batch_size,\n",
@@ -624,13 +551,20 @@
    "outputs": [],
    "source": [
     "# test and cache\n",
-    "dm = imdbHSDataModule(model, tokenizer, n=600)\n",
+    "dm = imdbHSDataModule(model, tokenizer, n=dataset_n, batch_size=batch_size, extract_layers=extract_layers)\n",
     "dm.setup('train')\n",
     "dl = dm.val_dataloader()\n",
     "b = next(iter(dl))\n",
     "b"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -647,7 +581,7 @@
    "outputs": [],
    "source": [
     "# test and cache\n",
-    "dm2 = imdbHSDataModule(model, tokenizer, prompt_fn=format_imdbs_multishot_lie, n=200)\n",
+    "dm2 = imdbHSDataModule(model, tokenizer, prompt_fn=format_imdbs_multishot_lie, n=dataset_n//6, batch_size=batch_size, extract_layers=extract_layers)\n",
     "dm2.setup('train')"
    ]
   },
@@ -688,9 +622,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "# plt.scatter(y, all_pos_ans)"
-   ]
+   "source": []
   },
   {
    "cell_type": "code",
@@ -1028,6 +960,173 @@
     "df_hist[['val/loss', 'train/loss']].plot()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## QC: Try a single pass"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_text_pairs = [\n",
+    "    # text, sentiment\n",
+    "    ['This movie was trash burger. It was a very bad movie.', 0],\n",
+    "    [\"This movie changed my life, I've watched it over 5 times and shown it to my entire family\", 1],\n",
+    "    [\"\"\"Lifetime did it again. Can we say stupid? I couldn't wait for it to end. The plot was senseless. The acting was terrible! Especially by the teenagers. The story has been played a thousand times! Are we just desperate to give actors a job? The previews were attractive and I was really looking for a good thriller.Once in awhile lifetime comes up with a good movie, this isn't one of them. Unless one has nothing else to do I would avoid this one at all cost. This was a waste of two hours of my life. Can I get them back? I would have rather scraped my face against a brick wall for two hours then soaked it in peroxide. That would have been more entertaining.\"\"\", 0],\n",
+    "    [\"I can't remember many films where a bumbling idiot of a hero was so funny throughout. Leslie Cheung is such the antithesis of a hero that he's too dense to be seduced by a gorgeous vampire... I had the good luck to see it on a big screen, and to find a video to watch again and again. 9/10\", 1],\n",
+    "    [\"The little girl Desi is so adorable... I cant think of a more beautiful story then this one here. It will make you cry, laugh, and believe. Knowing that this was based on a true story just made me gasp and it also made me realize that there are nice people out there. Great cast and an overall great movie.\", 1],    \n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# attempt at meta example....\n",
+    "\n",
+    "## Params\n",
+    "lie=0\n",
+    "question=0\n",
+    "\n",
+    "i = 3\n",
+    "text = [test_text_pairs[i][0]]\n",
+    "answer = test_text_pairs[i][1]\n",
+    "\n",
+    "## run\n",
+    "neg = get_hidden_states(model, tokenizer, format_imdbs_multishot(text, 0, lie=lie))\n",
+    "pos = get_hidden_states(model, tokenizer, format_imdbs_multishot(text, 1, lie=lie))\n",
+    "\n",
+    "hs = get_hidden_states(model, tokenizer, format_imdbs_multishot(text, question, lie=lie))\n",
+    "\n",
+    "## display\n",
+    "print(hs['text_q'][0])\n",
+    "print('='*80)\n",
+    "desired_ans=(question==answer)^lie\n",
+    "print(f\"question=q={question}, answer=a={answer}, lie=l={lie}. (q*a)^l==(({question}*{answer})^{lie}=={desired_ans}) \")\n",
+    "print(f'[public textual answer should be `{\"Yes\" if (question==answer)^lie  else \"No\"}` for this to be a {\"lie\" if lie else \"truth\"}:]')\n",
+    "print(hs['text_ans'][0])\n",
+    "print(f'[public numeric answer should be {\">50%\" if (desired_ans) else \"<50%\"}')\n",
+    "print(f\"{hs['ans'][0]:2.2%}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# FIXME also try with model \n",
+    "\n",
+    "neg = get_hidden_states(model, tokenizer, format_imdbs_multishot(text, 0, lie=lie))\n",
+    "pos = get_hidden_states(model, tokenizer, format_imdbs_multishot(text, 1, lie=lie))\n",
+    "b = 1\n",
+    "x0 = torch.from_numpy(neg['hidden_states']).reshape((b,-1)).float()#.unsqueeze(0)\n",
+    "x1 = torch.from_numpy(pos['hidden_states']).reshape((b,-1)).float()#.unsqueeze(0)\n",
+    "\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    batch = x0, x1, answer\n",
+    "    o = net.prediction_step(batch, 0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": null,