diff --git a/example.ipynb b/example.ipynb
index 202e8e1..4cc35de 100644
--- a/example.ipynb
+++ b/example.ipynb
@@ -37,19 +37,8 @@
      "text": [
       "/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/.venv/lib/python3.9/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
       "  warnings.warn(\n",
-      "`flash-attention` package not found, consider installing for better performance: No module named 'flash_attn'.\n",
-      "Current `flash-attenton` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`.\n",
-      "Loading checkpoint shards: 100%|██████████| 4/4 [00:02<00:00,  1.94it/s]\n",
       "/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/.venv/lib/python3.9/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
-      "  warnings.warn(\n",
-      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Loaded model and tokenizer\n"
+      "  warnings.warn(\n"
      ]
     }
    ],
@@ -58,8 +47,8 @@
     "import torch\n",
     "\n",
     "print(\"Loading model and tokenizer...\")\n",
-    "# model_name = \"databricks/dolly-v2-3b\"\n",
-    "model_name = \"failspy/kappa-3-phi-abliterated\"\n",
+    "model_name = \"databricks/dolly-v2-3b\"\n",
+    "# model_name = \"failspy/kappa-3-phi-abliterated\"\n",
     "model = AutoModelForCausalLM.from_pretrained(\n",
     "    model_name,\n",
     "    use_cache=True,\n",
@@ -67,135 +56,12 @@
     "    #  device=\"cuda:0\",\n",
     "    #  device_map=\"auto\",\n",
     "    attn_implementation='eager',\n",
-    "    trust_remote_code=True,\n",
+    "    # trust_remote_code=True,\n",
     ").to(\"cuda:0\")\n",
     "tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, use_cache=True)\n",
     "print(\"Loaded model and tokenizer\")"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Scratch"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# from jaxtyping import Float, Int\n",
-    "# import torch\n",
-    "# from torch.nn import functional as F\n",
-    "# from torch import Tensor\n",
-    "# from typing import List, Callable, Tuple, Dict, Optional\n",
-    "# import pandas as pd"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# # initital state\n",
-    "# prompt = \"The necromancer in his tower, what's his top problem? \"\n",
-    "# choices = [\"1\", \"the skeleton\", \"the boney boys\"]\n",
-    "# choices_tokens = tokenizer(choices).input_ids\n",
-    "# choices_tokens = [torch.tensor(c) for c in choices_tokens]\n",
-    "# # current_tokens = torch.tensor([])\n",
-    "\n",
-    "# # next\n",
-    "# input_ids = tokenizer([prompt], return_tensors=\"pt\").to(model.device).input_ids[0]\n",
-    "# choices_tokens\n",
-    "\n",
-    "# # for each next choice, continue down the tree, recording the log probs"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# def get_valid_next_choices(choices_tokens, current_tokens):\n",
-    "#     next_choices = []\n",
-    "#     for choice_tokens in choices_tokens:\n",
-    "#         # if we have some more slots left\n",
-    "#         if len(current_tokens)<len(choice_tokens):\n",
-    "#             # see if current_tokens matches\n",
-    "#             if (choice_tokens[: len(current_tokens)] == current_tokens).all():\n",
-    "#                 c = choice_tokens[len(current_tokens)].item()\n",
-    "#                 next_choices.append(c)\n",
-    "\n",
-    "#     next_choices = list(set(next_choices))\n",
-    "#     return torch.LongTensor(next_choices)\n",
-    "\n",
-    "\n",
-    "# def gen_choice_probs(\n",
-    "#     model: AutoModelForCausalLM,\n",
-    "#     tokenizer: AutoTokenizer,\n",
-    "#     input_ids: Int[Tensor, \"seq\"],\n",
-    "#     choices_tokens: List[Int[Tensor, \"seq\"]],\n",
-    "#     choice: Optional[Int[Tensor, \"\"]] = None,\n",
-    "#     prob: float = 1,\n",
-    "#     current_tokens: Int[Tensor, \"seq\"] = torch.LongTensor([]),\n",
-    "#     z=[],\n",
-    "# ):\n",
-    "#     if choice is not None:\n",
-    "#         c = choice[None].to(current_tokens.device)\n",
-    "#         current_tokens = torch.cat([current_tokens, c], dim=-1)\n",
-    "#         print(current_tokens, 'current_tokens')\n",
-    "#         c = choice[None].to(input_ids.device)\n",
-    "#         input_ids = torch.cat([input_ids, c], dim=-1)\n",
-    "\n",
-    "#     next_choices = get_valid_next_choices(choices_tokens, current_tokens)\n",
-    "#     if len(next_choices) == 0:\n",
-    "#         s = tokenizer.decode(current_tokens)\n",
-    "#         r = dict(tokens=current_tokens.cpu(), prob=prob, choice=s)\n",
-    "#         yield r\n",
-    "#     else:\n",
-    "#         o = model(input_ids[None])\n",
-    "#         logits_constrained = o.logits[0, -1][next_choices]\n",
-    "#         probs = F.softmax(logits_constrained, dim=-1)\n",
-    "#         for i in range(len(next_choices)):\n",
-    "#             next_choice = next_choices[i]\n",
-    "#             next_prob = prob * probs[i].item()\n",
-    "#             yield from gen_choice_probs(\n",
-    "#                 model=model,\n",
-    "#                 tokenizer=tokenizer,\n",
-    "#                 choices_tokens=choices_tokens,\n",
-    "#                 input_ids=input_ids,\n",
-    "#                 choice=next_choice,\n",
-    "#                 prob=next_prob,\n",
-    "#                 current_tokens=current_tokens,\n",
-    "#                 z=z + [i],\n",
-    "#             )\n",
-    "\n",
-    "\n",
-    "# r = list(gen_choice_probs(model, tokenizer, input_ids, choices_tokens))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# r"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# pd.DataFrame(r).sort_values(\"prob\", ascending=False).drop(columns=[\"tokens\"])"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -205,24 +71,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Generating...\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "You are not running the flash-attention implementation, expect numerical differences.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from prob_jsonformer.format import highlight_values\n",
     "from prob_jsonformer.main import Jsonformer\n",
@@ -235,7 +86,7 @@
     "            \"properties\": {\n",
     "                \"name\": {\"type\": \"string\"},\n",
     "                \"location\": {\"type\": \"string\"},\n",
-    "                \"choices\": {\"type\": \"choices\", \"enum\": [\"1\", \"the\", \"they walked the old dog\", \"1 the\"]},\n",
+    "                \"choice_probs\": {\"type\": \"choice_probs\", \"enum\": [\"ski\", \"snowboard\", \"walk\", \"pretend\"]},\n",
     "                \"inventory\": {\n",
     "                    \"type\": \"array\",\n",
     "                    \"items\": {\n",
@@ -245,9 +96,9 @@
     "                            \"name\": {\"type\": \"string\"},\n",
     "                            \"description\": {\"type\": \"string\"},\n",
     "                            \"category\": {\"type\": \"string\"},\n",
-    "                            # \"price\": {\"type\": \"number\"},\n",
+    "                            \"price\": {\"type\": \"number\"},\n",
     "                            \"inStock\": {\"type\": \"boolean\"},\n",
-    "                            # \"rating\": {\"type\": \"number\"},\n",
+    "                            \"rating\": {\"type\": \"number\"},\n",
     "                            \"images\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}},\n",
     "                        },\n",
     "                    },\n",
@@ -282,7 +133,9 @@
     "    \"type\": \"object\",\n",
     "    \"properties\": {\n",
     "        \"make\": {\"type\": \"string\"},\n",
-    "        \"model\": {\"type\": \"string\"},\n",
+    "        \"model\": {\"type\": \"choice_probs\", \"enum\": [\"Mazda\", \"Kea\"]},\n",
+    "        \"new\": {\"type\": \"choice_probs\", \"enum\": [\"true\", \"false\"]},\n",
+    "        \"rating\": {\"type\": \"choice_probs\", \"enum\": [\"1\", \"2\", \"3\", \"4\"]},\n",
     "        \"year\": {\"type\": \"number\"},\n",
     "        \"colors_available\": {\n",
     "            \"type\": \"array\",\n",
@@ -319,7 +172,9 @@
     "                \"make\": {\"type\": \"string\"},\n",
     "                \"model\": {\"type\": \"string\"},\n",
     "                \"year\": {\"type\": \"number\"},\n",
-    "                \"colors\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}},\n",
+    "                \"colors\": {\"type\": \"choice_probs\", \"enum\": [\"red\", \"green\", \"blue\", \"black\", \"white\"]},\n",
+    "                \"as_new\": {\"type\": \"choice_probs\", \"enum\": [\"true\", \"false\"]},\n",
+    "                \"rating\": {\"type\": \"choice_probs\", \"enum\": [\"1\", \"2\", \"3\", \"4\"]},\n",
     "                \"features\": {\n",
     "                    \"type\": \"object\",\n",
     "                    \"properties\": {\n",
diff --git a/prob_jsonformer/logits_processors.py b/prob_jsonformer/logits_processors.py
index db288d3..eafac18 100644
--- a/prob_jsonformer/logits_processors.py
+++ b/prob_jsonformer/logits_processors.py
@@ -2,6 +2,7 @@ from typing import List
 from transformers import PreTrainedTokenizer, LogitsWarper, StoppingCriteria
 import torch
 
+
 class StringStoppingCriteria(StoppingCriteria):
     def __init__(self, tokenizer: PreTrainedTokenizer, prompt_length: int):
         self.tokenizer = tokenizer
@@ -61,6 +62,7 @@ class NumberStoppingCriteria(StoppingCriteria):
 
         return False
 
+
 class OutputNumbersTokens(LogitsWarper):
     def __init__(self, tokenizer: PreTrainedTokenizer, prompt: str):
         self.tokenizer = tokenizer
@@ -78,7 +80,7 @@ class OutputNumbersTokens(LogitsWarper):
                 self.allowed_mask[token_id] = True
 
     def __call__(self, _, scores):
-        mask = self.allowed_mask.expand_as(scores)
+        mask = self.allowed_mask.expand_as(scores[0])
         scores[~mask] = -float("inf")
 
         return scores
diff --git a/prob_jsonformer/main.py b/prob_jsonformer/main.py
index f000d92..961530f 100644
--- a/prob_jsonformer/main.py
+++ b/prob_jsonformer/main.py
@@ -192,7 +192,7 @@ class Jsonformer:
             else:
                 obj.append(self.generation_marker)
             return self.generate_string()
-        elif schema_type == "choices":
+        elif schema_type == "choice_probs":
             if key:
                 obj[key] = self.generation_marker
             else: