diff --git a/CCS.ipynb b/.ipynb_checkpoints/001_mjc_CCS-checkpoint.ipynb
similarity index 52%
rename from CCS.ipynb
rename to .ipynb_checkpoints/001_mjc_CCS-checkpoint.ipynb
index 2242901..e9eacd5 100644
--- a/CCS.ipynb
+++ b/.ipynb_checkpoints/001_mjc_CCS-checkpoint.ipynb
@@ -11,42 +11,19 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "metadata": {},
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:25.921309Z",
+     "start_time": "2023-05-07T05:39:24.474456Z"
+    }
+   },
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/accounts/projects/jsteinhardt/uid1693600/.local/lib/python3.7/site-packages/pandas/compat/_optional.py:138: UserWarning: Pandas requires version '2.7.0' or newer of 'numexpr' (version '2.6.9' currently installed).\n",
-      "  warnings.warn(msg, UserWarning)\n",
-      "Reusing dataset amazon_polarity (/scratch/users/uid1693600/huggingface-cache/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc)\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "53469fc43b4542ab9fb1fa375896f9d3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/2 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
-      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
-      "Some weights of the model checkpoint at microsoft/deberta-v2-xxlarge were not used when initializing DebertaV2ForMaskedLM: ['lm_predictions.lm_head.dense.bias', 'lm_predictions.lm_head.LayerNorm.bias', 'lm_predictions.lm_head.LayerNorm.weight', 'deberta.embeddings.position_embeddings.weight', 'lm_predictions.lm_head.dense.weight', 'lm_predictions.lm_head.bias']\n",
-      "- This IS expected if you are initializing DebertaV2ForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
-      "- This IS NOT expected if you are initializing DebertaV2ForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
-      "Some weights of DebertaV2ForMaskedLM were not initialized from the model checkpoint at microsoft/deberta-v2-xxlarge and are newly initialized: ['cls.predictions.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight']\n",
-      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
      ]
     }
    ],
@@ -60,38 +37,243 @@
     "\n",
     "from datasets import load_dataset\n",
     "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM\n",
-    "from sklearn.linear_model import LogisticRegression\n",
-    "\n",
+    "from sklearn.linear_model import LogisticRegression"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "ExecuteTime": {
+     "start_time": "2023-05-07T01:08:20.635Z"
+    }
+   },
+   "source": [
+    "## Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:28.746274Z",
+     "start_time": "2023-05-07T05:39:25.922561Z"
+    },
+    "scrolled": false
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Found cached dataset amazon_polarity (/home/wassname/.cache/huggingface/datasets/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc)\n",
+      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 55.80it/s]\n"
+     ]
+    }
+   ],
+   "source": [
     "# Let's just try IMDB for simplicity\n",
-    "data = load_dataset(\"amazon_polarity\")[\"test\"]\n",
-    "\n",
+    "data = load_dataset(\"amazon_polarity\")[\"test\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def format_imdb(text, label):\n",
+    "    \"\"\"\n",
+    "    Given an imdb example (\"text\") and corresponding label (0 for negative, or 1 for positive), \n",
+    "    returns a zero-shot prompt for that example (which includes that label as the answer).\n",
+    "    \n",
+    "    (This is just one example of a simple, manually created prompt.)\n",
+    "    \"\"\"\n",
+    "    return \"The following movie review expresses a \" + [\"negative\", \"positive\"][label] + \" sentiment:\\n\" + text\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T03:55:24.085897Z",
+     "start_time": "2023-05-07T03:55:24.083858Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:37.745946Z",
+     "start_time": "2023-05-07T05:39:28.748845Z"
+    }
+   },
+   "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/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda117.so\n",
+      "CUDA_SETUP: WARNING! libcudart.so not found in any environmental path. Searching in backup paths...\n",
+      "CUDA SETUP: CUDA runtime path found: /usr/local/cuda/lib64/libcudart.so\n",
+      "CUDA SETUP: Highest compute capability among GPUs detected: 7.5\n",
+      "CUDA SETUP: Detected CUDA version 117\n",
+      "CUDA SETUP: Loading binary /home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda117.so...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: /home/wassname/miniforge3/envs/jupyter2 did not contain ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] as expected! Searching further paths...\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/usr/share/gconf/cinnamon.mandatory.path')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/usr/share/gconf/cinnamon.default.path')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('local/wassname-fractal-desktop'), PosixPath('@/tmp/.ICE-unix/5335,unix/wassname-fractal-desktop')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('0'), PosixPath('1')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/etc/xdg/xdg-cinnamon')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('module'), PosixPath('//matplotlib_inline.backend_inline')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: Found duplicate ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] files: {PosixPath('/usr/local/cuda/lib64/libcudart.so'), PosixPath('/usr/local/cuda/lib64/libcudart.so.11.0')}.. 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",
+      "Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:06<00:00,  3.24s/it]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "LlamaForCausalLM(\n",
+       "  (model): LlamaModel(\n",
+       "    (embed_tokens): Embedding(32000, 4096, padding_idx=0)\n",
+       "    (layers): ModuleList(\n",
+       "      (0-31): 32 x LlamaDecoderLayer(\n",
+       "        (self_attn): LlamaAttention(\n",
+       "          (q_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
+       "          (k_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
+       "          (v_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
+       "          (o_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
+       "          (rotary_emb): LlamaRotaryEmbedding()\n",
+       "        )\n",
+       "        (mlp): LlamaMLP(\n",
+       "          (gate_proj): Linear8bitLt(in_features=4096, out_features=11008, bias=False)\n",
+       "          (down_proj): Linear8bitLt(in_features=11008, out_features=4096, bias=False)\n",
+       "          (up_proj): Linear8bitLt(in_features=4096, out_features=11008, bias=False)\n",
+       "          (act_fn): SiLUActivation()\n",
+       "        )\n",
+       "        (input_layernorm): LlamaRMSNorm()\n",
+       "        (post_attention_layernorm): LlamaRMSNorm()\n",
+       "      )\n",
+       "    )\n",
+       "    (norm): LlamaRMSNorm()\n",
+       "  )\n",
+       "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
+       ")"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
     "# Here are a few different model options you can play around with:\n",
     "model_name = \"deberta\"\n",
-    "# model_name = \"gpt-j\"\n",
+    "model_name = \"gpt-j\"\n",
     "# model_name = \"t5\"\n",
+    "model_name = \"llama\"\n",
+    "# model_name = \"alpaca\"\n",
+    "finetuned = None\n",
+    "\n",
+    "model_options = dict(\n",
+    "    device_map=\"auto\", \n",
+    "    load_in_8bit=True,\n",
+    "    torch_dtype=torch.float16,\n",
+    ")\n",
     "\n",
-    "# if you want to cache the model weights somewhere, you can specify that here\n",
-    "cache_dir = None\n",
     "\n",
     "if model_name == \"deberta\":\n",
     "    model_type = \"encoder\"\n",
-    "    tokenizer = AutoTokenizer.from_pretrained(\"microsoft/deberta-v2-xxlarge\", cache_dir=cache_dir)\n",
-    "    model = AutoModelForMaskedLM.from_pretrained(\"microsoft/deberta-v2-xxlarge\", cache_dir=cache_dir)\n",
-    "    model.cuda()\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(\"microsoft/deberta-v2-xxlarge\")\n",
+    "    model = AutoModelForMaskedLM.from_pretrained(\"microsoft/deberta-v2-xxlarge\", **model_options)\n",
     "elif model_name == \"gpt-j\":\n",
     "    model_type = \"decoder\"\n",
-    "    tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6B\", cache_dir=cache_dir)\n",
-    "    model = AutoModelForCausalLM.from_pretrained(\"EleutherAI/gpt-j-6B\", cache_dir=cache_dir)\n",
-    "    model.cuda()\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6B\")\n",
+    "    model = AutoModelForCausalLM.from_pretrained(\"EleutherAI/gpt-j-6B\", **model_options)\n",
     "elif model_name == \"t5\":\n",
     "    model_type = \"encoder_decoder\"\n",
-    "    tokenizer = AutoTokenizer.from_pretrained(\"t5-11b\", cache_dir=cache_dir)\n",
-    "    model = AutoModelForSeq2SeqLM.from_pretrained(\"t5-11b\", cache_dir=cache_dir)\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(\"t5-11b\")\n",
+    "    model = AutoModelForSeq2SeqLM.from_pretrained(\"t5-11b\", **model_options)\n",
     "    model.parallelize()  # T5 is big enough that we may need to run it on multiple GPUs\n",
+    "elif (\"llama\" in model_name) or (\"alpaca\" in model_name):\n",
+    "    model_repo = \"Neko-Institute-of-Science/LLaMA-7B-HF\"\n",
+    "    lora_repo = \"tloen/alpaca-lora-7b\"\n",
+    "    model_type = \"decoder\"\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(model_repo)\n",
+    "    model = AutoModelForCausalLM.from_pretrained(model_repo, **model_options)\n",
+    "    \n",
+    "    if \"alpaca\" in model_name:\n",
+    "        from peft import PeftModel\n",
+    "        model = PeftModel.from_pretrained(\n",
+    "            model, \n",
+    "            lora_repo, \n",
+    "            device_map='auto'#{'': 0}\n",
+    "        )\n",
     "else:\n",
-    "    print(\"Not implemented!\")"
+    "    raise NotADirectoryError(model_name)\n",
+    "model"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T03:59:37.080463Z",
+     "start_time": "2023-05-07T03:59:37.074408Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T01:39:14.495263Z",
+     "start_time": "2023-05-07T01:39:14.495255Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -102,8 +284,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
+   "execution_count": 4,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:37.752425Z",
+     "start_time": "2023-05-07T05:39:37.747545Z"
+    }
+   },
    "outputs": [],
    "source": [
     "def get_encoder_hidden_states(model, tokenizer, input_text, layer=-1):\n",
@@ -174,6 +361,86 @@
     "    return fn(model, tokenizer, input_text, layer=layer)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:53.181466Z",
+     "start_time": "2023-05-07T05:39:53.179289Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# UPTO, fix nan"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:54.330398Z",
+     "start_time": "2023-05-07T05:39:53.506750Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([ 3.707e+00, -2.578e-01,  4.321e-01, ...,  8.389e-01,  2.798e-01,\n",
+       "       -1.083e-03], dtype=float16)"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# unit test\n",
+    "idx = 0\n",
+    "text, true_label = data[idx][\"content\"], data[idx][\"label\"]\n",
+    "neg_hs = get_hidden_states(model, tokenizer, format_imdb(text, 0), model_type=model_type)\n",
+    "neg_hs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:40:25.237818Z",
+     "start_time": "2023-05-07T05:40:25.235783Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# # sceatch\n",
+    "# layer = -10\n",
+    "# input_ids = tokenizer(text + tokenizer.eos_token, return_tensors=\"pt\").input_ids.to(model.device)\n",
+    "# # forward pass\n",
+    "# with torch.no_grad():\n",
+    "#     output = model(input_ids, output_hidden_states=True)\n",
+    "\n",
+    "# # get the last layer, last token hidden states\n",
+    "# hs_tuple = output[\"hidden_states\"]\n",
+    "# hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
+    "# hs, output['logits'], output['hidden_states']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:55.699656Z",
+     "start_time": "2023-05-07T05:39:55.653120Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -183,18 +450,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
+   "execution_count": 13,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:40:26.960648Z",
+     "start_time": "2023-05-07T05:40:26.954566Z"
+    }
+   },
    "outputs": [],
    "source": [
-    "def format_imdb(text, label):\n",
-    "    \"\"\"\n",
-    "    Given an imdb example (\"text\") and corresponding label (0 for negative, or 1 for positive), \n",
-    "    returns a zero-shot prompt for that example (which includes that label as the answer).\n",
-    "    \n",
-    "    (This is just one example of a simple, manually created prompt.)\n",
-    "    \"\"\"\n",
-    "    return \"The following movie review expresses a \" + [\"negative\", \"positive\"][label] + \" sentiment:\\n\" + text\n",
     "\n",
     "\n",
     "def get_hidden_states_many_examples(model, tokenizer, data, model_type, n=100):\n",
@@ -238,14 +502,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
+   "execution_count": 14,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:40:58.173815Z",
+     "start_time": "2023-05-07T05:40:27.095416Z"
+    }
+   },
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:19<00:00,  5.13it/s]\n"
+      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:31<00:00,  3.22it/s]\n"
      ]
     }
    ],
@@ -253,25 +522,45 @@
     "neg_hs, pos_hs, y = get_hidden_states_many_examples(model, tokenizer, data, model_type)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T03:15:48.077547Z",
+     "start_time": "2023-05-07T03:15:48.074666Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## 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 logistic regression accuracy is bad, there's no hope of CCS doing well."
+    "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."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
+   "execution_count": 23,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T06:49:28.092748Z",
+     "start_time": "2023-05-07T06:49:28.057699Z"
+    }
+   },
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Logistic regression accuracy: 0.92\n"
+      "Logistic regression accuracy: 1.0 [TRAIN]\n",
+      "Logistic regression accuracy: 0.8 [TEST]\n"
      ]
     }
    ],
@@ -289,9 +578,22 @@
     "\n",
     "lr = LogisticRegression(class_weight=\"balanced\")\n",
     "lr.fit(x_train, y_train)\n",
-    "print(\"Logistic regression accuracy: {}\".format(lr.score(x_test, y_test)))"
+    "print(\"Logistic regression accuracy: {} [TRAIN]\".format(lr.score(x_train, y_train)))\n",
+    "print(\"Logistic regression accuracy: {} [TEST]\".format(lr.score(x_test, y_test)))"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T06:48:49.955305Z",
+     "start_time": "2023-05-07T06:48:49.943711Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -301,20 +603,29 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
+   "execution_count": 24,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T07:32:53.940126Z",
+     "start_time": "2023-05-07T07:32:53.911010Z"
+    }
+   },
    "outputs": [],
    "source": [
     "class MLPProbe(nn.Module):\n",
     "    def __init__(self, d):\n",
     "        super().__init__()\n",
-    "        self.linear1 = nn.Linear(d, 100)\n",
-    "        self.linear2 = nn.Linear(100, 1)\n",
+    "        self.net = nn.Sequential(\n",
+    "            nn.Linear(d, 100),\n",
+    "            nn.ReLU(),\n",
+    "            nn.Linear(100, 100),\n",
+    "            nn.ReLU(),\n",
+    "            nn.Linear(100, 1),\n",
+    "            nn.Sigmoid(),\n",
+    "        )\n",
     "\n",
     "    def forward(self, x):\n",
-    "        h = F.relu(self.linear1(x))\n",
-    "        o = self.linear2(h)\n",
-    "        return torch.sigmoid(o)\n",
+    "        return torch.net(x)\n",
     "\n",
     "class CCS(object):\n",
     "    def __init__(self, x0, x1, nepochs=1000, ntries=10, lr=1e-3, batch_size=-1, \n",
@@ -440,18 +751,26 @@
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "CCS accuracy: 0.9\n"
-     ]
+   "cell_type": "markdown",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:40:58.250804Z",
+     "start_time": "2023-05-07T05:40:58.230537Z"
     }
-   ],
+   },
+   "source": [
+    "## Train"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "start_time": "2023-05-07T07:32:54.693Z"
+    }
+   },
+   "outputs": [],
    "source": [
     "# Train CCS without any labels\n",
     "ccs = CCS(neg_hs_train, pos_hs_train)\n",
@@ -462,6 +781,20 @@
     "print(\"CCS accuracy: {}\".format(ccs_acc))"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -472,9 +805,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "dlk2",
    "language": "python",
-   "name": "python3"
+   "name": "dlk2"
   },
   "language_info": {
    "codemirror_mode": {
@@ -486,7 +819,20 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.5"
+   "version": "3.9.16"
+  },
+  "toc": {
+   "base_numbering": 1,
+   "nav_menu": {},
+   "number_sections": true,
+   "sideBar": true,
+   "skip_h1_title": false,
+   "title_cell": "Table of Contents",
+   "title_sidebar": "Contents",
+   "toc_cell": false,
+   "toc_position": {},
+   "toc_section_display": true,
+   "toc_window_display": true
   },
   "vscode": {
    "interpreter": {
diff --git a/001_mjc_CCS.ipynb b/001_mjc_CCS.ipynb
new file mode 100644
index 0000000..e9eacd5
--- /dev/null
+++ b/001_mjc_CCS.ipynb
@@ -0,0 +1,845 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Let's implement CCS from scratch.\n",
+    "This will deliberately be a simple (but less efficient) implementation to make everything as clear as possible."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:25.921309Z",
+     "start_time": "2023-05-07T05:39:24.474456Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "from tqdm import tqdm\n",
+    "import copy\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "from datasets import load_dataset\n",
+    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM\n",
+    "from sklearn.linear_model import LogisticRegression"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "ExecuteTime": {
+     "start_time": "2023-05-07T01:08:20.635Z"
+    }
+   },
+   "source": [
+    "## Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:28.746274Z",
+     "start_time": "2023-05-07T05:39:25.922561Z"
+    },
+    "scrolled": false
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Found cached dataset amazon_polarity (/home/wassname/.cache/huggingface/datasets/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc)\n",
+      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 55.80it/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Let's just try IMDB for simplicity\n",
+    "data = load_dataset(\"amazon_polarity\")[\"test\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def format_imdb(text, label):\n",
+    "    \"\"\"\n",
+    "    Given an imdb example (\"text\") and corresponding label (0 for negative, or 1 for positive), \n",
+    "    returns a zero-shot prompt for that example (which includes that label as the answer).\n",
+    "    \n",
+    "    (This is just one example of a simple, manually created prompt.)\n",
+    "    \"\"\"\n",
+    "    return \"The following movie review expresses a \" + [\"negative\", \"positive\"][label] + \" sentiment:\\n\" + text\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T03:55:24.085897Z",
+     "start_time": "2023-05-07T03:55:24.083858Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:37.745946Z",
+     "start_time": "2023-05-07T05:39:28.748845Z"
+    }
+   },
+   "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/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda117.so\n",
+      "CUDA_SETUP: WARNING! libcudart.so not found in any environmental path. Searching in backup paths...\n",
+      "CUDA SETUP: CUDA runtime path found: /usr/local/cuda/lib64/libcudart.so\n",
+      "CUDA SETUP: Highest compute capability among GPUs detected: 7.5\n",
+      "CUDA SETUP: Detected CUDA version 117\n",
+      "CUDA SETUP: Loading binary /home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda117.so...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: /home/wassname/miniforge3/envs/jupyter2 did not contain ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] as expected! Searching further paths...\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/usr/share/gconf/cinnamon.mandatory.path')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/usr/share/gconf/cinnamon.default.path')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('local/wassname-fractal-desktop'), PosixPath('@/tmp/.ICE-unix/5335,unix/wassname-fractal-desktop')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('0'), PosixPath('1')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/etc/xdg/xdg-cinnamon')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('module'), PosixPath('//matplotlib_inline.backend_inline')}\n",
+      "  warn(msg)\n",
+      "/home/wassname/miniforge3/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: Found duplicate ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] files: {PosixPath('/usr/local/cuda/lib64/libcudart.so'), PosixPath('/usr/local/cuda/lib64/libcudart.so.11.0')}.. 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",
+      "Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:06<00:00,  3.24s/it]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "LlamaForCausalLM(\n",
+       "  (model): LlamaModel(\n",
+       "    (embed_tokens): Embedding(32000, 4096, padding_idx=0)\n",
+       "    (layers): ModuleList(\n",
+       "      (0-31): 32 x LlamaDecoderLayer(\n",
+       "        (self_attn): LlamaAttention(\n",
+       "          (q_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
+       "          (k_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
+       "          (v_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
+       "          (o_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
+       "          (rotary_emb): LlamaRotaryEmbedding()\n",
+       "        )\n",
+       "        (mlp): LlamaMLP(\n",
+       "          (gate_proj): Linear8bitLt(in_features=4096, out_features=11008, bias=False)\n",
+       "          (down_proj): Linear8bitLt(in_features=11008, out_features=4096, bias=False)\n",
+       "          (up_proj): Linear8bitLt(in_features=4096, out_features=11008, bias=False)\n",
+       "          (act_fn): SiLUActivation()\n",
+       "        )\n",
+       "        (input_layernorm): LlamaRMSNorm()\n",
+       "        (post_attention_layernorm): LlamaRMSNorm()\n",
+       "      )\n",
+       "    )\n",
+       "    (norm): LlamaRMSNorm()\n",
+       "  )\n",
+       "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
+       ")"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Here are a few different model options you can play around with:\n",
+    "model_name = \"deberta\"\n",
+    "model_name = \"gpt-j\"\n",
+    "# model_name = \"t5\"\n",
+    "model_name = \"llama\"\n",
+    "# model_name = \"alpaca\"\n",
+    "finetuned = None\n",
+    "\n",
+    "model_options = dict(\n",
+    "    device_map=\"auto\", \n",
+    "    load_in_8bit=True,\n",
+    "    torch_dtype=torch.float16,\n",
+    ")\n",
+    "\n",
+    "\n",
+    "if model_name == \"deberta\":\n",
+    "    model_type = \"encoder\"\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(\"microsoft/deberta-v2-xxlarge\")\n",
+    "    model = AutoModelForMaskedLM.from_pretrained(\"microsoft/deberta-v2-xxlarge\", **model_options)\n",
+    "elif model_name == \"gpt-j\":\n",
+    "    model_type = \"decoder\"\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6B\")\n",
+    "    model = AutoModelForCausalLM.from_pretrained(\"EleutherAI/gpt-j-6B\", **model_options)\n",
+    "elif model_name == \"t5\":\n",
+    "    model_type = \"encoder_decoder\"\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(\"t5-11b\")\n",
+    "    model = AutoModelForSeq2SeqLM.from_pretrained(\"t5-11b\", **model_options)\n",
+    "    model.parallelize()  # T5 is big enough that we may need to run it on multiple GPUs\n",
+    "elif (\"llama\" in model_name) or (\"alpaca\" in model_name):\n",
+    "    model_repo = \"Neko-Institute-of-Science/LLaMA-7B-HF\"\n",
+    "    lora_repo = \"tloen/alpaca-lora-7b\"\n",
+    "    model_type = \"decoder\"\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(model_repo)\n",
+    "    model = AutoModelForCausalLM.from_pretrained(model_repo, **model_options)\n",
+    "    \n",
+    "    if \"alpaca\" in model_name:\n",
+    "        from peft import PeftModel\n",
+    "        model = PeftModel.from_pretrained(\n",
+    "            model, \n",
+    "            lora_repo, \n",
+    "            device_map='auto'#{'': 0}\n",
+    "        )\n",
+    "else:\n",
+    "    raise NotADirectoryError(model_name)\n",
+    "model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T03:59:37.080463Z",
+     "start_time": "2023-05-07T03:59:37.074408Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T01:39:14.495263Z",
+     "start_time": "2023-05-07T01:39:14.495255Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## First let's write code for extracting hidden states given a model and text. \n",
+    "How we do this exactly will depend on the type of model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:37.752425Z",
+     "start_time": "2023-05-07T05:39:37.747545Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "def get_encoder_hidden_states(model, tokenizer, input_text, layer=-1):\n",
+    "    \"\"\"\n",
+    "    Given an encoder model and some text, gets the encoder hidden states (in a given layer, by default the last) \n",
+    "    on that input text (where the full text is given to the encoder).\n",
+    "\n",
+    "    Returns a numpy array of shape (hidden_dim,)\n",
+    "    \"\"\"\n",
+    "    # tokenize\n",
+    "    encoder_text_ids = tokenizer(input_text, truncation=True, return_tensors=\"pt\").input_ids.to(model.device)\n",
+    "\n",
+    "    # forward pass\n",
+    "    with torch.no_grad():\n",
+    "        output = model(encoder_text_ids, output_hidden_states=True)\n",
+    "\n",
+    "    # get the appropriate hidden states\n",
+    "    hs_tuple = output[\"hidden_states\"]\n",
+    "    \n",
+    "    hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
+    "\n",
+    "    return hs\n",
+    "\n",
+    "def get_encoder_decoder_hidden_states(model, tokenizer, input_text, layer=-1):\n",
+    "    \"\"\"\n",
+    "    Given an encoder-decoder model and some text, gets the encoder hidden states (in a given layer, by default the last) \n",
+    "    on that input text (where the full text is given to the encoder).\n",
+    "\n",
+    "    Returns a numpy array of shape (hidden_dim,)\n",
+    "    \"\"\"\n",
+    "    # tokenize\n",
+    "    encoder_text_ids = tokenizer(input_text, return_tensors=\"pt\").input_ids.to(model.device)\n",
+    "    decoder_text_ids = tokenizer(\"\", return_tensors=\"pt\").input_ids.to(model.device)\n",
+    "\n",
+    "    # forward pass\n",
+    "    with torch.no_grad():\n",
+    "        output = model(encoder_text_ids, decoder_input_ids=decoder_text_ids, output_hidden_states=True)\n",
+    "\n",
+    "    # get the appropriate hidden states\n",
+    "    hs_tuple = output[\"encoder_hidden_states\"]\n",
+    "    hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
+    "\n",
+    "    return hs\n",
+    "\n",
+    "def get_decoder_hidden_states(model, tokenizer, input_text, layer=-1):\n",
+    "    \"\"\"\n",
+    "    Given a decoder model and some text, gets the hidden states (in a given layer, by default the last) on that input text\n",
+    "\n",
+    "    Returns a numpy array of shape (hidden_dim,)\n",
+    "    \"\"\"\n",
+    "    # tokenize (adding the EOS token this time)\n",
+    "    input_ids = tokenizer(input_text + tokenizer.eos_token, return_tensors=\"pt\").input_ids.to(model.device)\n",
+    "\n",
+    "    # forward pass\n",
+    "    with torch.no_grad():\n",
+    "        output = model(input_ids, output_hidden_states=True)\n",
+    "\n",
+    "    # get the last layer, last token hidden states\n",
+    "    hs_tuple = output[\"hidden_states\"]\n",
+    "    hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
+    "\n",
+    "    return hs\n",
+    "\n",
+    "def get_hidden_states(model, tokenizer, input_text, layer=-1, model_type=\"encoder\"):\n",
+    "    fn = {\"encoder\": get_encoder_hidden_states, \"encoder_decoder\": get_encoder_decoder_hidden_states,\n",
+    "          \"decoder\": get_decoder_hidden_states}[model_type]\n",
+    "\n",
+    "    return fn(model, tokenizer, input_text, layer=layer)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:53.181466Z",
+     "start_time": "2023-05-07T05:39:53.179289Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# UPTO, fix nan"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:54.330398Z",
+     "start_time": "2023-05-07T05:39:53.506750Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([ 3.707e+00, -2.578e-01,  4.321e-01, ...,  8.389e-01,  2.798e-01,\n",
+       "       -1.083e-03], dtype=float16)"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# unit test\n",
+    "idx = 0\n",
+    "text, true_label = data[idx][\"content\"], data[idx][\"label\"]\n",
+    "neg_hs = get_hidden_states(model, tokenizer, format_imdb(text, 0), model_type=model_type)\n",
+    "neg_hs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:40:25.237818Z",
+     "start_time": "2023-05-07T05:40:25.235783Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# # sceatch\n",
+    "# layer = -10\n",
+    "# input_ids = tokenizer(text + tokenizer.eos_token, return_tensors=\"pt\").input_ids.to(model.device)\n",
+    "# # forward pass\n",
+    "# with torch.no_grad():\n",
+    "#     output = model(input_ids, output_hidden_states=True)\n",
+    "\n",
+    "# # get the last layer, last token hidden states\n",
+    "# hs_tuple = output[\"hidden_states\"]\n",
+    "# hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
+    "# hs, output['logits'], output['hidden_states']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:39:55.699656Z",
+     "start_time": "2023-05-07T05:39:55.653120Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Now let's write code for formatting data and for getting all the hidden states."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:40:26.960648Z",
+     "start_time": "2023-05-07T05:40:26.954566Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "def get_hidden_states_many_examples(model, tokenizer, data, model_type, n=100):\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",
+    "    all_neg_hs, all_pos_hs, all_gt_labels = [], [], []\n",
+    "\n",
+    "    # loop\n",
+    "    for _ in tqdm(range(n)):\n",
+    "        # for simplicity, sample a random example until we find one that's a reasonable length\n",
+    "        # (most examples should be a reasonable length, so this is just to make sure)\n",
+    "        while True:\n",
+    "            idx = np.random.randint(len(data))\n",
+    "            text, true_label = data[idx][\"content\"], data[idx][\"label\"]\n",
+    "            # the actual formatted input will be longer, so include a bit of a marign\n",
+    "            if len(tokenizer(text)) < 400:  \n",
+    "                break\n",
+    "                \n",
+    "        # get hidden states\n",
+    "        neg_hs = get_hidden_states(model, tokenizer, format_imdb(text, 0), model_type=model_type)\n",
+    "        pos_hs = get_hidden_states(model, tokenizer, format_imdb(text, 1), model_type=model_type)\n",
+    "\n",
+    "        # collect\n",
+    "        all_neg_hs.append(neg_hs)\n",
+    "        all_pos_hs.append(pos_hs)\n",
+    "        all_gt_labels.append(true_label)\n",
+    "\n",
+    "    all_neg_hs = np.stack(all_neg_hs)\n",
+    "    all_pos_hs = np.stack(all_pos_hs)\n",
+    "    all_gt_labels = np.stack(all_gt_labels)\n",
+    "\n",
+    "    return all_neg_hs, all_pos_hs, all_gt_labels"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:40:58.173815Z",
+     "start_time": "2023-05-07T05:40:27.095416Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:31<00:00,  3.22it/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "neg_hs, pos_hs, y = get_hidden_states_many_examples(model, tokenizer, data, model_type)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T03:15:48.077547Z",
+     "start_time": "2023-05-07T03:15:48.074666Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 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."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T06:49:28.092748Z",
+     "start_time": "2023-05-07T06:49:28.057699Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Logistic regression accuracy: 1.0 [TRAIN]\n",
+      "Logistic regression accuracy: 0.8 [TEST]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# let's create a simple 50/50 train split (the data is already randomized)\n",
+    "n = len(y)\n",
+    "neg_hs_train, neg_hs_test = neg_hs[:n//2], neg_hs[n//2:]\n",
+    "pos_hs_train, pos_hs_test = pos_hs[:n//2], pos_hs[n//2:]\n",
+    "y_train, y_test = y[:n//2], y[n//2:]\n",
+    "\n",
+    "# for simplicity we can just take the difference between positive and negative hidden states\n",
+    "# (concatenating also works fine)\n",
+    "x_train = neg_hs_train - pos_hs_train\n",
+    "x_test = neg_hs_test - pos_hs_test\n",
+    "\n",
+    "lr = LogisticRegression(class_weight=\"balanced\")\n",
+    "lr.fit(x_train, y_train)\n",
+    "print(\"Logistic regression accuracy: {} [TRAIN]\".format(lr.score(x_train, y_train)))\n",
+    "print(\"Logistic regression accuracy: {} [TEST]\".format(lr.score(x_test, y_test)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T06:48:49.955305Z",
+     "start_time": "2023-05-07T06:48:49.943711Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Now let's try CCS"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T07:32:53.940126Z",
+     "start_time": "2023-05-07T07:32:53.911010Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "class MLPProbe(nn.Module):\n",
+    "    def __init__(self, d):\n",
+    "        super().__init__()\n",
+    "        self.net = nn.Sequential(\n",
+    "            nn.Linear(d, 100),\n",
+    "            nn.ReLU(),\n",
+    "            nn.Linear(100, 100),\n",
+    "            nn.ReLU(),\n",
+    "            nn.Linear(100, 1),\n",
+    "            nn.Sigmoid(),\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return torch.net(x)\n",
+    "\n",
+    "class CCS(object):\n",
+    "    def __init__(self, x0, x1, nepochs=1000, ntries=10, lr=1e-3, batch_size=-1, \n",
+    "                 verbose=False, device=\"cuda\", linear=True, weight_decay=0.01, var_normalize=False):\n",
+    "        # data\n",
+    "        self.var_normalize = var_normalize\n",
+    "        self.x0 = self.normalize(x0)\n",
+    "        self.x1 = self.normalize(x1)\n",
+    "        self.d = self.x0.shape[-1]\n",
+    "\n",
+    "        # training\n",
+    "        self.nepochs = nepochs\n",
+    "        self.ntries = ntries\n",
+    "        self.lr = lr\n",
+    "        self.verbose = verbose\n",
+    "        self.device = device\n",
+    "        self.batch_size = batch_size\n",
+    "        self.weight_decay = weight_decay\n",
+    "        \n",
+    "        # probe\n",
+    "        self.linear = linear\n",
+    "        self.probe = self.initialize_probe()\n",
+    "        self.best_probe = copy.deepcopy(self.probe)\n",
+    "\n",
+    "        \n",
+    "    def initialize_probe(self):\n",
+    "        if self.linear:\n",
+    "            self.probe = nn.Sequential(nn.Linear(self.d, 1), nn.Sigmoid())\n",
+    "        else:\n",
+    "            self.probe = MLPProbe(self.d)\n",
+    "        self.probe.to(self.device)    \n",
+    "\n",
+    "\n",
+    "    def normalize(self, x):\n",
+    "        \"\"\"\n",
+    "        Mean-normalizes the data x (of shape (n, d))\n",
+    "        If self.var_normalize, also divides by the standard deviation\n",
+    "        \"\"\"\n",
+    "        normalized_x = x - x.mean(axis=0, keepdims=True)\n",
+    "        if self.var_normalize:\n",
+    "            normalized_x /= normalized_x.std(axis=0, keepdims=True)\n",
+    "\n",
+    "        return normalized_x\n",
+    "\n",
+    "        \n",
+    "    def get_tensor_data(self):\n",
+    "        \"\"\"\n",
+    "        Returns x0, x1 as appropriate tensors (rather than np arrays)\n",
+    "        \"\"\"\n",
+    "        x0 = torch.tensor(self.x0, dtype=torch.float, requires_grad=False, device=self.device)\n",
+    "        x1 = torch.tensor(self.x1, dtype=torch.float, requires_grad=False, device=self.device)\n",
+    "        return x0, x1\n",
+    "    \n",
+    "\n",
+    "    def get_loss(self, p0, p1):\n",
+    "        \"\"\"\n",
+    "        Returns the CCS loss for two probabilities each of shape (n,1) or (n,)\n",
+    "        \"\"\"\n",
+    "        informative_loss = (torch.min(p0, p1)**2).mean(0)\n",
+    "        consistent_loss = ((p0 - (1-p1))**2).mean(0)\n",
+    "        return informative_loss + consistent_loss\n",
+    "\n",
+    "\n",
+    "    def get_acc(self, x0_test, x1_test, y_test):\n",
+    "        \"\"\"\n",
+    "        Computes accuracy for the current parameters on the given test inputs\n",
+    "        \"\"\"\n",
+    "        x0 = torch.tensor(self.normalize(x0_test), dtype=torch.float, requires_grad=False, device=self.device)\n",
+    "        x1 = torch.tensor(self.normalize(x1_test), dtype=torch.float, requires_grad=False, device=self.device)\n",
+    "        with torch.no_grad():\n",
+    "            p0, p1 = self.best_probe(x0), self.best_probe(x1)\n",
+    "        avg_confidence = 0.5*(p0 + (1-p1))\n",
+    "        predictions = (avg_confidence.detach().cpu().numpy() < 0.5).astype(int)[:, 0]\n",
+    "        acc = (predictions == y_test).mean()\n",
+    "        acc = max(acc, 1 - acc)\n",
+    "\n",
+    "        return acc\n",
+    "    \n",
+    "        \n",
+    "    def train(self):\n",
+    "        \"\"\"\n",
+    "        Does a single training run of nepochs epochs\n",
+    "        \"\"\"\n",
+    "        x0, x1 = self.get_tensor_data()\n",
+    "        permutation = torch.randperm(len(x0))\n",
+    "        x0, x1 = x0[permutation], x1[permutation]\n",
+    "        \n",
+    "        # set up optimizer\n",
+    "        optimizer = torch.optim.AdamW(self.probe.parameters(), lr=self.lr, weight_decay=self.weight_decay)\n",
+    "        \n",
+    "        batch_size = len(x0) if self.batch_size == -1 else self.batch_size\n",
+    "        nbatches = len(x0) // batch_size\n",
+    "\n",
+    "        # Start training (full batch)\n",
+    "        for epoch in range(self.nepochs):\n",
+    "            for j in range(nbatches):\n",
+    "                x0_batch = x0[j*batch_size:(j+1)*batch_size]\n",
+    "                x1_batch = x1[j*batch_size:(j+1)*batch_size]\n",
+    "            \n",
+    "                # probe\n",
+    "                p0, p1 = self.probe(x0_batch), self.probe(x1_batch)\n",
+    "\n",
+    "                # get the corresponding loss\n",
+    "                loss = self.get_loss(p0, p1)\n",
+    "\n",
+    "                # update the parameters\n",
+    "                optimizer.zero_grad()\n",
+    "                loss.backward()\n",
+    "                optimizer.step()\n",
+    "\n",
+    "        return loss.detach().cpu().item()\n",
+    "    \n",
+    "    def repeated_train(self):\n",
+    "        best_loss = np.inf\n",
+    "        for train_num in range(self.ntries):\n",
+    "            self.initialize_probe()\n",
+    "            loss = self.train()\n",
+    "            if loss < best_loss:\n",
+    "                self.best_probe = copy.deepcopy(self.probe)\n",
+    "                best_loss = loss\n",
+    "\n",
+    "        return best_loss"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-05-07T05:40:58.250804Z",
+     "start_time": "2023-05-07T05:40:58.230537Z"
+    }
+   },
+   "source": [
+    "## Train"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "start_time": "2023-05-07T07:32:54.693Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# Train CCS without any labels\n",
+    "ccs = CCS(neg_hs_train, pos_hs_train)\n",
+    "ccs.repeated_train()\n",
+    "\n",
+    "# Evaluate\n",
+    "ccs_acc = ccs.get_acc(neg_hs_test, pos_hs_test, y_test)\n",
+    "print(\"CCS accuracy: {}\".format(ccs_acc))"
+   ]
+  },
+  {
+   "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": "dlk2"
+  },
+  "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"
+  },
+  "toc": {
+   "base_numbering": 1,
+   "nav_menu": {},
+   "number_sections": true,
+   "sideBar": true,
+   "skip_h1_title": false,
+   "title_cell": "Table of Contents",
+   "title_sidebar": "Contents",
+   "toc_cell": false,
+   "toc_position": {},
+   "toc_section_display": true,
+   "toc_window_display": true
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "b80286374679f2ad472c61c83fc267d31329b5dea8e2dcaccb727123767724c5"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/mjc_notes.md b/mjc_notes.md
new file mode 100644
index 0000000..7dfbe3e
--- /dev/null
+++ b/mjc_notes.md
@@ -0,0 +1,7 @@
+```sh
+conda create -n dlk2 python=3.9 -y
+conda activate dlk2
+mamba install -y pytorch torchvision torchaudio pytorch-cuda=11.7 cudatoolkit-dev==11.7  cudatoolkit=11.7 -c pytorch -c nvidia  -c conda-forge
+mamba install -y ipykernel  pip
+pip install -r requirements.txt
+```
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..2a68d44
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,8 @@
+datasets
+promptsource
+tqdm
+transformers
+sklearn
+scikit-learn
+accelerate
+bitsandbytes