diff --git a/notebooks/017_mjc_sup_mcdrop_dm.ipynb b/notebooks/017_mjc_sup_mcdrop_dm.ipynb
new file mode 100644
index 0000000..5febc79
--- /dev/null
+++ b/notebooks/017_mjc_sup_mcdrop_dm.ipynb
@@ -0,0 +1,1696 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Lets just do supervised learning\n",
+    "\n",
+    "Since we are looking at pairs with random permuations (from dropout), we can't use CCS. This is because our probabilities do not add to one.\n",
+    "\n",
+    "People question if unsupervised learning bings anything to the table anyway, so lets start with supervised..."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "links:\n",
+    "- [loading](https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/alpaca.py)\n",
+    "- [dict](https://github.com/deep-diver/LLM-As-Chatbot/blob/c79e855a492a968b54bac223e66dc9db448d6eba/model_cards.json#L143)\n",
+    "- [prompt_format](https://github.com/deep-diver/PingPong/blob/main/src/pingpong/alpaca.py)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "import copy\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "from matplotlib import pyplot as plt\n",
+    "plt.style.use('ggplot')\n",
+    "\n",
+    "import random\n",
+    "from typing import Optional, List, Dict, Union\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch import Tensor\n",
+    "from torch import optim\n",
+    "from torch.utils.data import random_split, DataLoader, TensorDataset\n",
+    "\n",
+    "import pickle\n",
+    "import hashlib\n",
+    "from pathlib import Path\n",
+    "\n",
+    "from datasets import load_dataset\n",
+    "import datasets\n",
+    "\n",
+    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM, AutoConfig\n",
+    "import transformers\n",
+    "from transformers.models.auto.modeling_auto import AutoModel\n",
+    "from transformers import LogitsProcessorList\n",
+    "\n",
+    "\n",
+    "import lightning.pytorch as pl\n",
+    "from dataclasses import dataclass\n",
+    "\n",
+    "from sklearn.linear_model import LogisticRegression\n",
+    "# from scipy.stats import zscore\n",
+    "from sklearn.metrics import f1_score, roc_auc_score, accuracy_score\n",
+    "from sklearn.preprocessing import RobustScaler\n",
+    "\n",
+    "from tqdm.auto import tqdm\n",
+    "import gc\n",
+    "import os\n",
+    "\n",
+    "from loguru import logger\n",
+    "logger.add(os.sys.stderr, format=\"{time} {level} {message}\", level=\"INFO\")\n",
+    "\n",
+    "\n",
+    "transformers.__version__"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Model\n",
+    "\n",
+    "Chosing:\n",
+    "- https://old.reddit.com/r/LocalLLaMA/wiki/models\n",
+    "- https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard\n",
+    "- https://github.com/deep-diver/LLM-As-Chatbot/blob/main/model_cards.json\n",
+    "\n",
+    "\n",
+    "A uncensored and large one might be best for lying."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from peft import PeftModel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# leaderboard https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard\n",
+    "model_options = dict(\n",
+    "    device_map=\"auto\", \n",
+    "    load_in_4bit=True,\n",
+    "    # load_in_8bit=True,\n",
+    "    torch_dtype=torch.float16,\n",
+    "    trust_remote_code=True,\n",
+    "    use_safetensors=False,\n",
+    "    # use_cache=False,\n",
+    ")\n",
+    "\n",
+    "# so I need to use either pythia, stablelm, or tiiuae/falcon-7b-instruct to get dropout...\n",
+    "# moel_repo = \"stabilityai/stablelm-tuned-alpha-7b\" # poor performance\n",
+    "\n",
+    "# https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/falcon.py\n",
+    "model_repo = \"tiiuae/falcon-7b-instruct\"\n",
+    "# model_repo = \"tiiuae/falcon-7b\"\n",
+    "# model_repo = \"togethercomputer/RedPajama-INCITE-7B-Instruct\"\n",
+    "# model_repo = \"OpenAssis/tant/oasst-sft-4-pythia-12b-epoch-3.5\"\n",
+    "# model_repo = \"OpenAssistant/falcon-7b-sft-top1-696\"\n",
+    "# model_repo = \"openaccess-ai-collective/manticore-13b\"\n",
+    "# model_repo = \"TheBloke/Wizard-Vicuna-13B-Uncensored-HF\"\n",
+    "# model_repo = \"dvruette/llama-13b-pretrained-dropout\"\n",
+    "# model_repo = \"elinas/llama-13b-hf-transformers-4.29\" # no dropout\n",
+    "# lora_repo = \"LLMs/AlpacaGPT4-LoRA-13B-elina\"\n",
+    "model_repo = \"bigcode/starcoderplus\"\n",
+    "model_repo = \"HuggingFaceH4/starchat-beta\"\n",
+    "model_repo = \"WizardLM/WizardCoder-15B-V1.0\"\n",
+    "# model_repo= \"~/.cache/huggingface/hub/models--HuggingFaceH4--starchat-beta\"\n",
+    "# lora_repo = None\n",
+    "lora_repo = None\n",
+    "\n",
+    "config = AutoConfig.from_pretrained(model_repo, trust_remote_code=True,)\n",
+    "print(config)\n",
+    "# config.attn_pdrop=0.3\n",
+    "# config.embd_pdrop=0.3\n",
+    "# config.resid_pdrop=0.3\n",
+    "config.use_cache = False\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_repo)\n",
+    "model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options)\n",
+    "\n",
+    "if lora_repo is not None:\n",
+    "    # https://github.com/tloen/alpaca-lora/blob/main/generate.py#L40\n",
+    "    from peft import PeftModel\n",
+    "    model = PeftModel.from_pretrained(\n",
+    "        model,\n",
+    "        lora_repo, \n",
+    "        torch_dtype=torch.float16,\n",
+    "        lora_dropout=0.2,\n",
+    "        device_map='auto'\n",
+    "    )\n",
+    "    \n",
+    "# if not mode_8bit and not mode_4bit:\n",
+    "#     model.half()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 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",
+    "    tokenizer.pad_token_id = 204 # <unk> https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/alpaca.py\n",
+    "tokenizer.padding_side = \"left\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Params"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Params\n",
+    "N_SAMPLES = 9000\n",
+    "BATCH_SIZE = 8 # None # None means auto # 6 gives 16Gb/25GB. where 10GB is the base model. so 6 is 6/15\n",
+    "N_SHOTS = 3\n",
+    "USE_MCDROPOUT = True\n",
+    "dataset_n = 200\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",
+    "extract_layers = tuple(range(2, num_layers, stride)) + (num_layers,)\n",
+    "extract_layers, num_layers"
+   ]
+  },
+  {
+   "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": [
+    "# 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": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Load Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Let's just try IMDB for simplicity\n",
+    "dataset = load_dataset(\"amazon_polarity\")\n",
+    "data = dataset['test']"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Define Prompt\n",
+    "\n",
+    "- Lillian Wang's guide https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\n",
+    "- Elvi's guide https://www.promptingguide.ai/techniques/fewshot\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "example_prompts = dataset['train'].shuffle(seed=42)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "def random_example():\n",
+    "    \"\"\"gets a random example from shuffled train and formats it a bit\"\"\"\n",
+    "    ex = dict(content=('blah blah '*302))\n",
+    "    while len(tokenizer(ex['content']).input_ids) > 150: # 95% are below this 250, 70% below 150\n",
+    "        i = np.random.randint(example_prompts.num_rows)\n",
+    "        ex = example_prompts[i]\n",
+    "    input = f\"Title: \\\"{ex['title']}\\\". Content: \\\"{ex['content']}\\\"\"\n",
+    "    return input, ex['label']==1\n",
+    "\n",
+    "print(random_example()[0])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prefix_lie = prefix_true = prefix = f\"\"\"The classic puzzle of the two guards. It presents a scenario where two guards are involved in a conflict of truths and lies. Here's the setup:\n",
+    "\n",
+    "There are two guards standing at two separate doors. One door leads to safety, while the other leads to danger. One guard always tells the truth, and the other guard always lies. You don't know which guard is which or which door leads to safety.\"\"\"\n",
+    "\n",
+    "char_lie = \"Left Guard\"\n",
+    "char_true = \"Left Guard\"\n",
+    "\n",
+    "def prompt_format_alpaca(input:str, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n",
+    "    \"\"\"alpaca format\"\"\"\n",
+    "    prefix = \"\"\n",
+    "    if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n",
+    "    char = char_lie if lie else char_true\n",
+    "    instruction = f'Following the above instructions, using your own character, classify the sentiment of the given movie review, \"positive\" or \"negative\".'\n",
+    "    alpaca_prompt = f'{prefix}### Instruction:\\n{instruction}\\n\\n{input}\\n\\n### {char} Response:\\n{response}'\n",
+    "    return alpaca_prompt\n",
+    "\n",
+    "\n",
+    "def prompt_format_oa(input:str, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n",
+    "    \"\"\"alpaca format\"\"\"\n",
+    "    prefix = \"\"\n",
+    "    if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"<|endoftext|>\"\n",
+    "    char = char_lie if lie else char_true\n",
+    "    instruction = f'Following the above instructions, using your own character, classify the sentiment of the given movie review, \"positive\" or \"negative\".'\n",
+    "    if not response==\"\": response+=\"<|endoftext|>\"\n",
+    "    alpaca_prompt = f'{prefix}<|prompter|>{instruction}\\n{input}<|endoftext|><|assistant|>{char} Response:\\n{response}'\n",
+    "    return alpaca_prompt\n",
+    "\n",
+    "def prompt_format_falcon(input:str, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n",
+    "    prefix = \"\"\n",
+    "    if include_prefix: prefix = \"Instruction:\\n\" + (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n",
+    "    char = char_lie if lie else char_true\n",
+    "    instruction = f'Following the above instructions, using your own character, classify the sentiment of the given movie review, \"positive\" or \"negative\".'\n",
+    "    alpaca_prompt = f'{prefix}Question:\\n{instruction}\\n\\nContext:\\n{input}\\n\\nAnswer:\\n{response}'\n",
+    "    return alpaca_prompt\n",
+    "\n",
+    "\n",
+    "def prompt_format_vicuna(input:str, question:Optional[bool]=None, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n",
+    "    \"\"\"\n",
+    "    vicuna format\n",
+    "    \n",
+    "    https://github.com/melodysdreamj/WizardVicunaLM\n",
+    "    \"\"\"\n",
+    "    prefix = \"\"\n",
+    "    if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n",
+    "    instruction = f\"Is the sentiment of the below review {'positive' if (question==1) else 'negative'}?\"\n",
+    "    alpaca_prompt = f'{prefix}USER: {instruction} {input}\\nASSISTANT: {response}'\n",
+    "    return alpaca_prompt\n",
+    "\n",
+    "# def prompt_format_vicuna2(input:str, question:Optional[bool]=None, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n",
+    "#     \"\"\"\n",
+    "#     vicuna format\n",
+    "    \n",
+    "#     https://github.com/melodysdreamj/WizardVicunaLM\n",
+    "#     \"\"\"\n",
+    "#     prefix = \"\"\n",
+    "#     if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n",
+    "#     instruction = f\"Is the sentiment of the below review {'positive' if (question==1) else 'negative'}?\"\n",
+    "#     alpaca_prompt = f'{prefix}USER: {instruction} {input}\\nAssistant:\\n{response}'\n",
+    "#     return alpaca_prompt\n",
+    "\n",
+    "def prompt_format_manticore(input:str, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n",
+    "    \"\"\"\n",
+    "    vicuna format\n",
+    "    \n",
+    "    https://github.com/melodysdreamj/WizardVicunaLM\n",
+    "    https://huggingface.co/openaccess-ai-collective/manticore-13b#examples\n",
+    "    \"\"\"\n",
+    "    prefix = \"\"\n",
+    "    if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n",
+    "    char = char_lie if lie else char_true\n",
+    "    instruction = f'Classify the sentiment of the given movie review, \"positive\" or \"negative\".'\n",
+    "    alpaca_prompt = f'{prefix}### Instruction: {instruction}\\n\\n{input}\\n\\n### {char}:\\n{response}'\n",
+    "    return alpaca_prompt\n",
+    "\n",
+    "# def prompt_format_manticore2(input:str, question:Optional[bool]=None, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n",
+    "#     \"\"\"\n",
+    "#     vicuna format\n",
+    "    \n",
+    "#     https://github.com/melodysdreamj/WizardVicunaLM\n",
+    "#     https://huggingface.co/openaccess-ai-collective/manticore-13b#examples\n",
+    "#     \"\"\"\n",
+    "#     prefix = \"\"\n",
+    "#     if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n",
+    "#     instruction = f\"Is the sentiment of the below review {'positive' if (question==1) else 'negative'}?\"\n",
+    "#     alpaca_prompt = f'{prefix}USER: {instruction} {input}\\nASSISTANT: {response}'\n",
+    "#     return alpaca_prompt\n",
+    "\n",
+    "def prompt_format_chatml(input:str, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n",
+    "    \"\"\"\n",
+    "    https://huggingface.co/HuggingFaceH4/starchat-beta\n",
+    "    \n",
+    "    \"<|system|>\\n<|end|>\\n<|user|>\\n{query}<|end|>\\n<|assistant|>\"\n",
+    "    \"\"\"\n",
+    "    prefix = \"\"\n",
+    "    if include_prefix: prefix = \"<|system|>\" + (prefix_lie if lie else prefix_true) + \"<|end|>\\n\"\n",
+    "    char = char_lie if lie else char_true\n",
+    "    if len(response)>0:\n",
+    "        response += \"<|end|>\\n\"\n",
+    "    instruction = f'Classify the sentiment of the given movie review, \"positive\" or \"negative\".'\n",
+    "    alpaca_prompt = f'{prefix}<|user|>{instruction}\\n\\n{input}\\n\\n<|end|>\\n<|assistant|>\\n{response}'\n",
+    "    return alpaca_prompt\n",
+    "\n",
+    "\n",
+    "repo_dict = {\n",
+    "    \"TheBloke/Wizard-Vicuna-13B-Uncensored-HF\": 'vicuna',\n",
+    "    'Neko-Institute-of-Science/VicUnLocked-30b-LoRA': 'vicuna',\n",
+    "    \"ehartford/Wizard-Vicuna-13B-Uncensored\": 'vicuna',\n",
+    "    \"HuggingFaceH4/starchat-beta\": 'chatml',\n",
+    "    \"WizardLM/WizardCoder-15B-V1.0\": 'alpaca',\n",
+    "    # 'tiiuae/falcon-7b': 'manticore',\n",
+    "    # 'tiiuae/falcon-7b-instruct': 'vicuna',\n",
+    "}\n",
+    "prompt_formats = {\n",
+    "    'vicuna': prompt_format_vicuna,\n",
+    "    'alpaca': prompt_format_alpaca,\n",
+    "    'llama': prompt_format_alpaca,\n",
+    "    'manticore': prompt_format_manticore,\n",
+    "    'falcon': prompt_format_falcon,\n",
+    "    'chatml': prompt_format_chatml,\n",
+    "}\n",
+    "def guess_prompt_format(model_repo, lora_repo):\n",
+    "    repo = model_repo if (lora_repo is None) else lora_repo\n",
+    "    if repo in repo_dict:\n",
+    "        prompt_type = repo_dict[repo]\n",
+    "        return prompt_formats[prompt_type]\n",
+    "    for fmt in prompt_formats:\n",
+    "        if fmt in repo.lower():\n",
+    "            fn = prompt_formats[fmt]\n",
+    "            print(f\"guessing prompt format '{str(fn.__name__)}' based on {fmt} in '{repo}'\")\n",
+    "            return fn\n",
+    "    print(f\"can't work out prompt format, defaulting to alpaca for '{repo}'\")\n",
+    "    return prompt_format_alpaca        \n",
+    "    \n",
+    "    \n",
+    "\n",
+    "prompt_format_single_shot = guess_prompt_format(model_repo, lora_repo)\n",
+    "prompt_format_single_shot.__name__"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "def set_seeds(n):\n",
+    "    transformers.set_seed(n)\n",
+    "    torch.manual_seed(n)\n",
+    "    np.random.seed(n)\n",
+    "    random.seed(n)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rand_bool = lambda : np.random.rand()>0.5\n",
+    "\n",
+    "\n",
+    "def to_item(x):\n",
+    "    if isinstance(x, torch.Tensor):\n",
+    "        x = x.detach().cpu().item()\n",
+    "    return x\n",
+    "\n",
+    "\n",
+    "def format_imdb_multishot(input:str, response:str=\"\", lie:Optional[bool]=None, n_shots=N_SHOTS, verbose:bool=False, answer:Optional[bool]=None, seed=None):\n",
+    "    if seed is not None:\n",
+    "        set_seeds(seed)\n",
+    "    if lie is None: \n",
+    "        lie = rand_bool()\n",
+    "    if seed is not None:\n",
+    "        set_seeds(seed)\n",
+    "    \n",
+    "    main = prompt_format_single_shot(input, response, lie=lie)\n",
+    "    desired_answer = answer^lie == 1 if answer is not None else None\n",
+    "    info = dict(input=input, lie=lie, desired_answer=desired_answer, true_answer=answer)\n",
+    "    \n",
+    "    shots = []\n",
+    "    for i in range(n_shots):\n",
+    "        \n",
+    "        input, answer = random_example()\n",
+    "        # question=rand_bool()\n",
+    "        desired_answer = (answer)^lie == 1\n",
+    "        if verbose: print(f\"shot-{i} answer={answer}, lie={lie}. (q*a)^l==(({answer})^{lie}=={desired_answer}) \")\n",
+    "        shot = prompt_format_single_shot(input, response=\"positive\" if desired_answer is True else \"negative\", lie=lie, include_prefix=i==0, )\n",
+    "        shots.append(shot)\n",
+    "        \n",
+    "        \n",
+    "    info = {k:to_item(v) for k,v in info.items()}    \n",
+    "\n",
+    "    return \"\\n\\n\".join(shots+[main]), info\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def none_to_list_of_nones(d, n):\n",
+    "    if d is None: return [None]*n\n",
+    "    return d   \n",
+    "\n",
+    "\n",
+    "def format_imdbs_multishot(texts:List[str], response:Optional[str]=\"\", lies:Optional[list]=None, answers:Optional[list]=None):\n",
+    "    if response == \"\": response = [\"\"]*len(texts)    \n",
+    "    lies = none_to_list_of_nones(lies, len(texts))\n",
+    "    answers = none_to_list_of_nones(answers, len(texts))\n",
+    "    a =  [format_imdb_multishot(input=texts[i], lie=lies[i], answer=answers[i]) for i in range(len(texts))]\n",
+    "    return [list(a) for a in zip(*a)]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# QC: generation\n",
+    "\n",
+    "Let's a quick generation, so we can QC the output and sanity check that the model can actually do the task"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "text, label = random_example()\n",
+    "q, info = format_imdb_multishot(text, answer=label, lie=True, verbose=True)\n",
+    "print(q)\n",
+    "print('-'*80)\n",
+    "pipeline = transformers.pipeline(\n",
+    "    \"text-generation\",\n",
+    "    model=model,\n",
+    "    tokenizer=tokenizer,\n",
+    ")\n",
+    "sequences = pipeline(\n",
+    "    q,\n",
+    "    max_length=800,\n",
+    "    do_sample=False,\n",
+    "    return_full_text=False,\n",
+    "    eos_token_id=tokenizer.eos_token_id,\n",
+    ")\n",
+    "for seq in sequences:\n",
+    "    print(f\"{seq['generated_text']}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Guess batch size"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_size_dict = {\n",
+    "    \"HuggingFaceH4/starchat-beta\": '13b',\n",
+    "    'WizardLM/WizardCoder-15B-V1.0': '13b', # actually 15b\n",
+    "}\n",
+    "\n",
+    "\n",
+    "def guess_batch_size(model_repo, N_SHOTS):\n",
+    "    \"\"\"Some rougth guestimates of batch size. \n",
+    "    \n",
+    "    Aiming to undershoot rather than crash.\"\"\"\n",
+    "    if model_repo in model_size_dict:\n",
+    "        model_repo = model_size_dict[model_repo]\n",
+    "    \n",
+    "    if '7b' in model_repo.lower():\n",
+    "        return int(48//(2+N_SHOTS))\n",
+    "    elif '13b' in model_repo.lower():\n",
+    "        return int(24//(2+N_SHOTS))\n",
+    "    elif '30b' in model_repo.lower(): \n",
+    "        return int(6//(2+N_SHOTS))\n",
+    "    else:\n",
+    "        raise NotImplementedError(f\"can't work out size of '{model_repo}'\")\n",
+    "    \n",
+    "if BATCH_SIZE is None:\n",
+    "    BATCH_SIZE = guess_batch_size(model_repo, N_SHOTS)\n",
+    "    print(f\"guessing BATCH_SIZE {BATCH_SIZE} for '{model_repo}'\")\n",
+    "guess_batch_size('7b', N_SHOTS), guess_batch_size('13b', N_SHOTS), guess_batch_size('30b', N_SHOTS)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Collect hidden state pairs\n",
+    "\n",
+    "The idea is this: given two pairs of hidden states, where everything is the same except the random seed or dropout. Then tell me which one is more truthfull? \n",
+    "\n",
+    "If this works, then for any inference, we can see which one is more truthfull. Then we can see if it's the lower or higher probability one, and judge the answer and true or false.\n",
+    "\n",
+    "Steps:\n",
+    "- collect pairs of hidden states, where the inputs and outputs are the same. We modify the random seed and dropout.\n",
+    "- Each pair should have a binary answer. We can get that by comparing the probabilities of two tokens such as Yes and No.\n",
+    "- Train a prob to distinguish the pairs as more and less truthfull\n",
+    "- Test probe to see if it generalizes"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def clear_mem():\n",
+    "    gc.collect()\n",
+    "    torch.cuda.empty_cache()\n",
+    "    gc.collect()\n",
+    "    \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",
+    "                \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",
+    "                return True\n",
+    "    return False\n",
+    "    \n",
+    "clear_mem()\n",
+    "assert check_for_dropout(model), 'model should have dropout modules'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "            \n",
+    "def get_hidden_states(model, tokenizer, input_text, layers=extract_layers, truncation_length=900, 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",
+    "    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",
+    "        \n",
+    "    # Handling truncation: truncate start, not end\n",
+    "    if truncation_length is not None:\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",
+    "        \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",
+    "        \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.detach().cpu()[:, last_token] for v in attentions]\n",
+    "            attentions = torch.concat(attentions).numpy()\n",
+    "        \n",
+    "        hidden_states = torch.stack([outputs['hidden_states'][i] for i in layers], 1).detach().cpu().numpy()\n",
+    "        \n",
+    "        hidden_states = hidden_states[:, :, last_token] # (batch, layers, past_seq, logits) take just the last token so they are same size\n",
+    "        \n",
+    "        text_q = 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",
+    "\n",
+    "        scores = outputs['scores'][:, first_token].softmax(-1).detach().cpu().numpy() # 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",
+    "    return dict(hidden_states=hidden_states, ans=ans, text_ans=text_ans, text_q=text_q, input_id_shape=input_ids.shape,\n",
+    "                attentions=attentions, prob_n=prob_n, prob_y=prob_y, scores=outputs['scores'][:, 0].detach().cpu()\n",
+    "               )\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "clear_mem()"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Helper  Batch data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cache_dir = Path(\".pkl_cache\")\n",
+    "cache_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "def md5hash(s: str) -> str:\n",
+    "    return hashlib.md5(s).hexdigest()\n",
+    "\n",
+    "def cache_strargs_kwargs(func):\n",
+    "        \n",
+    "    def wrap(model, tokenizer, data, prompt_fn, n, batch_size):\n",
+    "        \"\"\"wrapper to cache results\"\"\"        \n",
+    "        # some args are to big (model), some are irrelavent (batch_size) and some the function name are not enougth (promt_fn)\n",
+    "        # so lets do some custom key to make sure we cache bust well\n",
+    "        set_seeds(42)\n",
+    "        text, label = random_example()\n",
+    "        example_prompt1, _ = format_imdb_multishot(text, answer=True, lie=True, seed=42)\n",
+    "        example_prompt2, _ = format_imdb_multishot(text, answer=False, lie=False, seed=42)\n",
+    "        kwargs = [str(model), str(tokenizer), str(data), str(prompt_fn.__name__), n, example_prompt1, example_prompt2,]\n",
+    "        logger.debug(f\"kwargs {kwargs}\")\n",
+    "        \n",
+    "        # The file name contains the hash of functions args and kwargs\n",
+    "        key = pickle.dumps(kwargs, 1)\n",
+    "        hsh = md5hash(key)[:6]\n",
+    "        f = cache_dir / f\"{hsh}.pkl\"\n",
+    "        \n",
+    "        \n",
+    "        if f.exists():\n",
+    "            logger.info(f\"loading hs from {f}\")\n",
+    "            res = pickle.load(f.open('rb'))\n",
+    "        else:\n",
+    "            res = func(model, tokenizer, data, prompt_fn, n, batch_size)\n",
+    "            logger.info(f\"caching hs to {f}\")\n",
+    "            pickle.dump(res, f.open('wb'))\n",
+    "        return res\n",
+    "    \n",
+    "    return wrap\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "@cache_strargs_kwargs\n",
+    "def batch_hidden_states(model, tokenizer, data, prompt_fn, n=100, batch_size=2):\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",
+    "    \n",
+    "    res = []\n",
+    "    infos = []\n",
+    "    \n",
+    "    ds_subset = data.shuffle(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",
+    "        texts, true_labels = batch[\"content\"], batch[\"label\"]\n",
+    "        lies = [i%2==0 for i,_ in enumerate(texts)] # every second one will be a lie\n",
+    "        q, info = format_imdbs_multishot(texts, answers=true_labels, lies=lies)\n",
+    "        if i==0:\n",
+    "            assert len(texts)==len(prompt_fn(texts, 0)[0]), 'make sure the prompt function can handle a list of text'\n",
+    "        \n",
+    "        \n",
+    "        # differen't due to dropout\n",
+    "        hs1 = get_hidden_states(model, tokenizer, q)\n",
+    "        hs2 = get_hidden_states(model, tokenizer, q)\n",
+    "        if i==0:\n",
+    "            assert hs1['hidden_states'][0, 0, 0]-hs2['hidden_states'][0, 0, 0]>0.001, \"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",
+    "        # collect\n",
+    "        b = len(texts)\n",
+    "        res.append([\n",
+    "            hs1['hidden_states'].reshape((b,-1)),\n",
+    "            hs1[\"ans\"],       \n",
+    "            hs2['hidden_states'].reshape((b,-1)),\n",
+    "            hs2[\"ans\"],\n",
+    "            true_labels,\n",
+    "        ])\n",
+    "        infos += info\n",
+    "        \n",
+    "        \n",
+    "        clear_mem()\n",
+    "    \n",
+    "    res = [np.concatenate(r) for r in zip(*res)]\n",
+    "    return *res, infos"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Lightning DataModule"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class imdbHSDataModule(pl.LightningDataModule):\n",
+    "\n",
+    "    def __init__(self,\n",
+    "                 model: AutoModel,\n",
+    "                 tokenizer: AutoTokenizer,\n",
+    "                 prompt_fn=format_imdbs_multishot,\n",
+    "                 dataset_name=\"amazon_polarity\",\n",
+    "                 batch_size=BATCH_SIZE,\n",
+    "                 dl_batch_size=32,\n",
+    "                 n=6000,\n",
+    "                ):\n",
+    "        super().__init__()\n",
+    "        self.save_hyperparameters(ignore=[\"model\", \"tokenizer\", \"prompt_fn\"])\n",
+    "        self.tokenizer = tokenizer\n",
+    "        self.model = model\n",
+    "        self.prompt_fn=prompt_fn\n",
+    "        \n",
+    "        self.dataset = None\n",
+    "\n",
+    "    def setup(self, stage: str):\n",
+    "        h = self.hparams\n",
+    "        \n",
+    "        # just setup once\n",
+    "        if self.dataset is not None:\n",
+    "            print('skipping setup, using cached values')\n",
+    "            return None\n",
+    "\n",
+    "        self.dataset = load_dataset(h.dataset_name, split=\"test\")\n",
+    "\n",
+    "        # in ELK they cache as a huggingface dataset\n",
+    "        self.hs1, self.ans1, self.hs2, self.ans2, self.y, self.infos = batch_hidden_states(\n",
+    "            self.model, self.tokenizer, self.dataset, self.prompt_fn, n=h.n, 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",
+    "        self.val_split = vs = int(n * 0.5)\n",
+    "        self.test_split = ts = int(n * 0.75)\n",
+    "        hs1_train, hs2_train, y_train = self.hs1[:vs], self.hs2[:vs], self.y[:vs]\n",
+    "        hs1_val, hs2_val, y_val = self.hs1[vs:ts], self.hs2[vs:ts], self.y[vs:ts]\n",
+    "        hs1_test, hs2_test, y_test = self.hs1[ts:],self. hs2[ts:], self.y[ts:]\n",
+    "        \n",
+    "        # make a dataframe for non hidden states\n",
+    "        self.df = pd.DataFrame(self.infos)\n",
+    "        self.df['ans1'] = self.ans1\n",
+    "        self.df['ans2'] = self.ans2\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 = hs1_train - hs2_train\n",
+    "        self.x_val = hs1_val - hs2_val\n",
+    "        self.x_test = hs1_test - hs2_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",
+    "        self.ds_train = TensorDataset(torch.from_numpy(hs1_train).float(),\n",
+    "                                      torch.from_numpy(hs2_train).float(),\n",
+    "                                      torch.from_numpy(y_train).float())\n",
+    "\n",
+    "        self.ds_val = TensorDataset(torch.from_numpy(hs1_val).float(),\n",
+    "                                    torch.from_numpy(hs2_val).float(),\n",
+    "                                    torch.from_numpy(y_val).float())\n",
+    "\n",
+    "        self.ds_test = TensorDataset(torch.from_numpy(hs1_test).float(),\n",
+    "                                     torch.from_numpy(hs2_test).float(),\n",
+    "                                     torch.from_numpy(y_test).float())\n",
+    "\n",
+    "    def train_dataloader(self):\n",
+    "        return DataLoader(self.ds_train,\n",
+    "                          batch_size=self.hparams.dl_batch_size,\n",
+    "                          shuffle=True)\n",
+    "\n",
+    "    def val_dataloader(self):\n",
+    "        return DataLoader(self.ds_val, batch_size=self.hparams.dl_batch_size)\n",
+    "\n",
+    "    def test_dataloader(self):\n",
+    "        return DataLoader(self.ds_test, batch_size=self.hparams.dl_batch_size)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# test and cache\n",
+    "dm = imdbHSDataModule(model, tokenizer, batch_size=BATCH_SIZE, n=N_SAMPLES)\n",
+    "dm.setup('train')\n",
+    "\n",
+    "dl_val = dm.val_dataloader()\n",
+    "dl_train = dm.train_dataloader()\n",
+    "b = next(iter(dl_train))\n",
+    "clear_mem()\n",
+    "b"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hss1 = dm.hs1\n",
+    "hss2 = dm.hs2\n",
+    "ans_1 = dm.ans1\n",
+    "ans_2 = dm.ans2\n",
+    "infos = dm.infos"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# temp: balance everything in case we stopped early\n",
+    "print(len(infos), len(ans_1), len(ans_2))\n",
+    "hss1 = hss1[:len(hss2)]\n",
+    "hss2 = hss2[:len(hss1)]\n",
+    "ans_1 = ans_1[:len(ans_2)]\n",
+    "ans_2 = ans_2[:len(ans_1)]\n",
+    "infos = infos[:len(ans_2)]\n",
+    "\n",
+    "df_infos2 = pd.DataFrame(infos)\n",
+    "df_infos2['dir_true'] = ans_2 - ans_1\n",
+    "df_infos2"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Model Results"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Task results\n",
+    "\n",
+    "E.g. how well does the underlying language model do on the task"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ans = (ans_1 + ans_2) / 2\n",
+    "acc=((ans>0.5)==df_infos2['true_answer']).mean()\n",
+    "print(f\"acc {acc:2.2f}\")\n",
+    "\n",
+    "d = df_infos2['lie']==True\n",
+    "acc = ((ans[d]>0.5)==df_infos2[d]['true_answer']).mean()\n",
+    "print(f\"acc when lie=True {acc:2.2f}\")\n",
+    "\n",
+    "d = df_infos2['lie']==False\n",
+    "acc = ((ans[d]>0.5)==df_infos2[d]['true_answer']).mean()\n",
+    "print(f\"acc when lie=False {acc:2.2f}\")\n",
+    "# ((ans_1>0)==df_infos2['desired_answer']).mean()"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Data prep\n",
+    "\n",
+    "We do two inferences on the same inputs. Since we have dropout enabled, even during inference, we get two slightly different hidden states `hs1` and `hs2`, and two slightly different probabilities for our yes and no output tokens `p1` `p2`. We also have the true answer `t`\n",
+    "\n",
+    "So there are a few ways we can set up the problem. \n",
+    "\n",
+    "We can vary x:\n",
+    "- `model(hs1)-model(hs2)=y`\n",
+    "- `model(hs1-hs2)==y`\n",
+    "\n",
+    "And we can try differen't y's:\n",
+    "- direction with a ranked loss. This could be unsupervised.\n",
+    "- magnitude with a regression loss\n",
+    "- vector (direction and magnitude) with a regression loss"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# QC: Linear supervised probes\n",
+    "\n",
+    "\n",
+    "Let's verify that the model's representations are good\n",
+    "\n",
+    "Before trying CCS, let's make sure there exists a direction that classifies examples as true vs false with high accuracy; if supervised logistic regression accuracy is bad, there's no hope of unsupervised CCS doing well.\n",
+    "\n",
+    "Note that because logistic regression is supervised we expect it to do better but to have worse generalisation that equivilent unsupervised methods. However in this case CSS is using a deeper model so it is more complicated.\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Try a classification of direction to truth"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "n = len(df_infos2)\n",
+    "\n",
+    "# Define X and y\n",
+    "X = hss1-hss2\n",
+    "\n",
+    "y = y_dir = df_infos2['true_answer'] == (df_infos2['dir_true']>0) # direction\n",
+    "\n",
+    "# split\n",
+    "n = len(y)\n",
+    "print('split size', n//2)\n",
+    "X_train, X_test = X[:n//2], X[n//2:]\n",
+    "y_train, y_test = y[:n//2], y[n//2:]\n",
+    "\n",
+    "# scale\n",
+    "scaler = RobustScaler()\n",
+    "scaler.fit(X_train)\n",
+    "X_train2 = scaler.transform(X_train)\n",
+    "X_test2 = scaler.transform(X_test)\n",
+    "\n",
+    "lr = LogisticRegression(class_weight=\"balanced\", penalty=\"l2\", max_iter=380)\n",
+    "lr.fit(X_train2, y_train>0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"Logistic cls acc: {:2.2%} [TRAIN]\".format(lr.score(X_train2, y_train>0)))\n",
+    "print(\"Logistic cls acc: {:2.2%} [TEST]\".format(lr.score(X_test2, y_test>0)))\n",
+    "\n",
+    "m = df_infos2['lie'][n//2:]\n",
+    "y_test_pred = lr.predict(X_test2)\n",
+    "acc_w_lie = ((y_test_pred[m]>0)==(y_test[m]>0)).mean()\n",
+    "acc_wo_lie = ((y_test_pred[~m]>0)==(y_test[~m]>0)).mean()\n",
+    "print(f'test acc w lie {acc_w_lie:2.2%}')\n",
+    "print(f'test acc wo lie {acc_wo_lie:2.2%}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_info_test = df_infos2.iloc[n//2:].copy()\n",
+    "y_pred = lr.predict(X_test2)\n",
+    "df_info_test['inner_truth'] = y_pred\n",
+    "df_info_test"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Result, detecting deception?"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "lie_pred = df_info_test['inner_truth']==df_info_test['true_answer']\n",
+    "lie_true = df_info_test['lie']\n",
+    "acc_lie = accuracy_score(lie_pred, lie_true)\n",
+    "print(f\"model can detect lies with acc {acc_lie:2.2%}\")\n",
+    "print(f\"w lies {sum(lie_true)}/{len(lie_true)} test rows\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Try a regression of the vector (magnitude and direction) vs truth"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bool_to_switch = lambda b:b*2-1\n",
+    "true_answer_switch = bool_to_switch(df_infos2['true_answer'])\n",
+    "y = y_left_more_true = df_infos2['dir_true'] * true_answer_switch\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Try a regression\n",
+    "from sklearn.linear_model import ElasticNet\n",
+    "\n",
+    "# Try a classification of direction\n",
+    "n = len(df_infos2)\n",
+    "\n",
+    "# Define X and y\n",
+    "X = hss1-hss2\n",
+    "y = y_left_more_true * 10\n",
+    "\n",
+    "# split\n",
+    "# y = df_infos2['dir2'] * 100\n",
+    "n = len(y)\n",
+    "print('split size', n//2)\n",
+    "X_train, X_test = X[:n//2], X[n//2:]\n",
+    "y_train, y_test = y[:n//2], y[n//2:]\n",
+    "\n",
+    "# scale\n",
+    "scaler = RobustScaler()\n",
+    "scaler.fit(X_train)\n",
+    "X_train2 = scaler.transform(X_train)\n",
+    "X_test2 = scaler.transform(X_test)\n",
+    "\n",
+    "X_train2 = X_train\n",
+    "X_test2 = X_test2\n",
+    "\n",
+    "lr2 = ElasticNet(max_iter=1000,)\n",
+    "lr2.fit(X_train2, y_train)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "eps = 0.\n",
+    "acc=np.mean((lr2.predict(X_train2)>eps)==(y_train>eps))\n",
+    "print(f'acc from train ElasticNet {acc:2.2f}')\n",
+    "acc=np.mean((lr2.predict(X_test2)>eps)==(y_test>eps))\n",
+    "print(f'acc from test ElasticNet {acc:2.2f}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "y_test_pred = lr2.predict(X_test)\n",
+    "plt.scatter(y_test, y_test_pred)\n",
+    "plt.xlabel('true')\n",
+    "plt.ylabel('pred')\n",
+    "plt.title('pred vs true on test')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# LightningModel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class MLPProbe(nn.Module):\n",
+    "    def __init__(self, d, depth=1, hs=32, dropout=0):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        layers = [\n",
+    "            nn.BatchNorm1d(d),  # this will normalise the inputs\n",
+    "            nn.Linear(d, hs),\n",
+    "            nn.Dropout1d(dropout),\n",
+    "        ]\n",
+    "        for _ in range(depth):\n",
+    "            layers += [\n",
+    "                nn.Linear(hs, hs),\n",
+    "                nn.ReLU(),\n",
+    "                nn.Dropout1d(dropout),\n",
+    "            ]\n",
+    "        layers += [nn.Linear(hs, 1)]\n",
+    "        self.net = nn.Sequential(*layers)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return self.net(x)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# logit0 = (torch.rand(5, 4)-0.5)*100\n",
+    "# logit1 = (torch.rand(5, 4)-0.5)*100\n",
+    "# ccs_squared_loss(logit0, logit1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pytorch_optimizer import Ranger21\n",
+    "import torchmetrics\n",
+    "\n",
+    "    \n",
+    "class CSS(pl.LightningModule):\n",
+    "    def __init__(self, d, total_steps, lr=4e-3, weight_decay=1e-9, dropout=0):\n",
+    "        super().__init__()\n",
+    "        self.probe = MLPProbe(d, depth=1, dropout=dropout)\n",
+    "        self.save_hyperparameters()\n",
+    "        self.auroc = torchmetrics.Accuracy(task=\"multiclass\", num_classes=2)\n",
+    "        \n",
+    "    def forward(self, x):\n",
+    "        return self.probe(x)\n",
+    "        \n",
+    "    def _step(self, batch, batch_idx, stage='train'):\n",
+    "        x0, x1, y = batch\n",
+    "        logit0, logit1 = self(x0), self(x1)\n",
+    "        logits = torch.concatenate([logit0, logit1], 1)\n",
+    "        y_pred = F.softmax(logits, -1)\n",
+    "        if stage=='pred':\n",
+    "            return y_pred\n",
+    "        \n",
+    "        loss = F.cross_entropy(logits, y.long())\n",
+    "        self.log(f\"{stage}/loss\", loss)\n",
+    "        \n",
+    "        self.auroc(y_pred, y.long())\n",
+    "        self.log(f\"{stage}/acc_step\", self.auroc)      \n",
+    "        return loss\n",
+    "    \n",
+    "    def on_train_epoch_end(self):\n",
+    "        # log epoch metric\n",
+    "        self.log('train/acc_epoch', self.auroc)\n",
+    "    \n",
+    "    def training_step(self, batch, batch_idx):\n",
+    "        return self._step(batch, batch_idx)\n",
+    "    \n",
+    "    def validation_step(self, batch, batch_idx=0):\n",
+    "        return self._step(batch, batch_idx, stage='val')\n",
+    "    \n",
+    "    def predict_step(self, batch, batch_idx):\n",
+    "        return self._step(batch, batch_idx, stage='pred')\n",
+    "\n",
+    "    # def configure_optimizers(self):\n",
+    "    #     optimizer = optim.AdamW(self.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay)\n",
+    "    #     lr_scheduler = optim.lr_scheduler.OneCycleLR(\n",
+    "    #         optimizer, self.hparams.lr, total_steps=self.hparams.total_steps\n",
+    "    #     )\n",
+    "        # return [optimizer], [lr_scheduler]\n",
+    "    \n",
+    "    def configure_optimizers(self):\n",
+    "        \"\"\"use ranger21 from  https://github.com/kozistr/pytorch_optimizer\"\"\"\n",
+    "        optimizer = Ranger21(\n",
+    "            self.parameters(),\n",
+    "            lr=self.hparams.lr,\n",
+    "            weight_decay=self.hparams.weight_decay,       \n",
+    "            num_iterations=self.hparams.total_steps,\n",
+    "        )\n",
+    "        return optimizer\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Run"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# quiet please\n",
+    "torch.set_float32_matmul_precision('medium')\n",
+    "\n",
+    "import warnings\n",
+    "warnings.filterwarnings(\"ignore\", \".*does not have many workers.*\")\n",
+    "warnings.filterwarnings(\"ignore\", \".*F-score.*\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Prep dataloader/set"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# split\n",
+    "X = hss1-hss2\n",
+    "y = (df_infos2['true_answer'] == (df_infos2['dir_true']>0)).values # direction\n",
+    "n = len(y)\n",
+    "print('split size', n//2)\n",
+    "\n",
+    "neg_hs_train = hss1[:n//2]\n",
+    "pos_hs_train = hss2[:n//2]\n",
+    "\n",
+    "neg_hs_val = hss1[n//2:]\n",
+    "pos_hs_val = hss2[n//2:]\n",
+    "\n",
+    "y_train, y_val = y[:n//2], y[n//2:]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dl_train = dm.train_dataloader()\n",
+    "dl_val = dm.val_dataloader()\n",
+    "b = next(iter(dl_train))\n",
+    "# b"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# init the model\n",
+    "max_epochs = 33\n",
+    "d = b[0].shape[-1]\n",
+    "net = CSS(d=d, total_steps=max_epochs*len(dl_train), lr=4e-3, weight_decay=1e-3, dropout=0.3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# with torch.no_grad():\n",
+    "#     b = next(iter(dl_train))\n",
+    "#     b2 = [bb.to(net.device) for bb in b]\n",
+    "#     y = net(b2[0])\n",
+    "# y"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # DEBUG\n",
+    "# trainer = pl.Trainer(fast_dev_run=2)\n",
+    "# trainer.fit(model=net, train_dataloaders=dl_train)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = pl.Trainer(\n",
+    "                     max_epochs=max_epochs, log_every_n_steps=5)\n",
+    "trainer.fit(model=net, train_dataloaders=dl_train, val_dataloaders=dl_val)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Read hist"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# import pytorch_lightning as pl\n",
+    "from lightning.pytorch.loggers.csv_logs import CSVLogger\n",
+    "# from pytorch_lightning.loggers.csv_logs import CSVLogger as CSVLogger2\n",
+    "from pathlib import Path\n",
+    "import pandas as pd\n",
+    "\n",
+    "def read_metrics_csv(metrics_file_path):\n",
+    "    df_hist = pd.read_csv(metrics_file_path)\n",
+    "    df_hist[\"epoch\"] = df_hist[\"epoch\"].ffill()\n",
+    "    df_histe = df_hist.set_index(\"epoch\").groupby(\"epoch\").mean()\n",
+    "    return df_histe\n",
+    "\n",
+    "\n",
+    "def read_hist(trainer: pl.Trainer):\n",
+    "\n",
+    "    ts = [t for t in trainer.loggers if isinstance(t, CSVLogger)]\n",
+    "    print(ts)\n",
+    "    try:\n",
+    "        metrics_file_path = Path(ts[0].experiment.metrics_file_path)\n",
+    "        df_histe = read_metrics_csv(metrics_file_path)\n",
+    "        return df_histe\n",
+    "    except Exception as e:\n",
+    "        raise e\n",
+    "        \n",
+    "        \n",
+    "df_hist = read_hist(trainer).ffill().bfill()\n",
+    "df_hist\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "keys = set(s.split('/')[1] for s in df_hist.columns if '/' in s)\n",
+    "for k in keys:    \n",
+    "    df_hist[[c for c in df_hist.columns if c.endswith(k)]].plot(title=k)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# df_hist[['val/acc', 'train/acc']].plot()\n",
+    "\n",
+    "# # df_hist[['val/f1', 'train/f1']].plot()\n",
+    "\n",
+    "# # df_hist[['val/roc_auc_bc', 'train/roc_auc_bc']].plot()\n",
+    "\n",
+    "# # df_hist[['val/roc_auc_mc', 'train/roc_auc_mc']].plot()\n",
+    "\n",
+    "# df_hist[['val/loss', 'train/loss']].plot()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Predict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dl_test = dm.test_dataloader()\n",
+    "y_test_pred = trainer.predict(net, dl_test)\n",
+    "y_test_pred = np.concatenate(y_test_pred)\n",
+    "# y_test_pred"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_test = dm.df.iloc[dm.val_split:dm.test_split].copy()\n",
+    "df_test['pred'] = y_test_pred.argmax(-1)\n",
+    "df_test"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "acc_truth = (df_test['pred']==df_test['true_answer']).mean()\n",
+    "print(f\"lightning model acc at predicting truth: {acc_truth:2.2%}\")"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "OK this doesn't work because for the CSS loss to work, it must add up to 1... ours does not"
+   ]
+  },
+  {
+   "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": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "dlk2",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.16"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}