diff --git a/README.md b/README.md
index fdb36b8..f7b342e 100644
--- a/README.md
+++ b/README.md
@@ -1,50 +1,67 @@
 # prob_jsonformer: A Bulletproof Way to Generate Probabilistic Structured JSON from Language Models.
 
-This fork has been modified to include the token probabilities. The original README is below.
+This fork has been modified to include the token probabilities. The original [README](https://github.com/1rgs/jsonformer) is includesbelow.
+
+I've also merged some of hte recent pr's for enum, integer, null, union. You can see them all below in this example:
 
 
 ## Example
 
 ```python
-from jsonformer import Jsonformer
+from prob_jsonformer import Jsonformer
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-12b")
-tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
+model_name = "databricks/dolly-v2-3b"
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 json_schema = {
     "type": "object",
     "properties": {
         "name": {"type": "string"},
-        "age": {"type": "choice_probs", "enum": ["8", "9", "10", "11"]},
-        "age2": {"type": "number"},
+        "age": {"type": "integer"},
+        "age_probs": {"type": "choice_probs", "enum": [str(s) for s in range(10, 20)]},
+        "unit_time": {"type": "number"},
         "is_student": {"type": "boolean"},
-        "is_student2": {"type": "choice_probs", "enum": ["true", "false"]},
+        "is_student_probs": {"type": "choice_probs", "enum": ["true", "false"]},
         "courses": {
             "type": "array",
             "items": {"type": "string"}
-        }
+        },
+        "trim": {"type": ["string", "null"]},
+        "color": {
+            "type": "enum",
+            "values": ["red", "green", "blue", "brown", "white", "black"],
+        },
     }
 }
 
-prompt = "Generate a person's information based on the following schema:"
-jsonformer = Jsonformer(model, tokenizer, json_schema, prompt)
+prompt = "Generate a young person's information based on the following schema:"
+jsonformer = Jsonformer(model, tokenizer, json_schema, prompt, temperature=0)
 generated_data = jsonformer()
 
-print(generated_data)
-# {'name': 'John Doe',
-#  'age': [{'prob': 0.1497802734375, 'choice': '8'},
-#   {'prob': 0.159423828125, 'choice': '9'},
-#   {'prob': 0.0982666015625, 'choice': '11'},
-#   {'prob': 0.59228515625, 'choice': '10'}],
-#  'age2': 10.0201,
-#  'is_student': True,
-#  'is_student2': [{'prob': 0.94580078125, 'choice': 'true'},
-#   {'prob': 0.05419921875, 'choice': 'false'}],
-#  'courses': ['C++']}
+generated_data = {'name': 'John Doe',
+ 'age': 20,
+ 'age_probs': [{'prob': 0.794921875, 'choice': '12'},
+  {'prob': 0.068359375, 'choice': '10'},
+  {'prob': 0.04345703125, 'choice': '16'},
+  {'prob': 0.03228759765625, 'choice': '14'},
+  {'prob': 0.0175628662109375, 'choice': '11'},
+  {'prob': 0.0157318115234375, 'choice': '15'},
+  {'prob': 0.006664276123046875, 'choice': '18'},
+  {'prob': 0.0046539306640625, 'choice': '13'},
+  {'prob': 0.00041294097900390625, 'choice': '17'},
+  {'prob': 0.00028824806213378906, 'choice': '19'},,
+ 'unit_time': 0.01,
+ 'is_student': True,
+ 'is_student_probs': [{'prob': 0.8310546875, 'choice': 'true'},
+  {'prob': 0.1688232421875, 'choice': 'false'}],
+ 'courses': ['C1'],
+ 'trim': None,
+ 'color': 'white'}
 ```
 
-# Jsonformer: A Bulletproof Way to Generate Structured JSON from Language Models.
+# ORIGINAL: Jsonformer: A Bulletproof Way to Generate Structured JSON from Language Models.
 
 ### Problem: Getting models to output structured JSON is hard
 
diff --git a/example.ipynb b/example.ipynb
index 5ccbc1d..7807ed8 100644
--- a/example.ipynb
+++ b/example.ipynb
@@ -111,12 +111,36 @@
       "        name: \u001b[32m\"Snowboard X-15\"\u001b[0m,\n",
       "        description: \u001b[32m\"Snowboard for all levels\"\u001b[0m,\n",
       "        category: \u001b[32m\"Snowboards\"\u001b[0m,\n",
-      "        price: \u001b[32m20.0375\u001b[0m,\n",
+      "        price: \u001b[32m20.0\u001b[0m,\n",
       "        inStock: \u001b[32mTrue\u001b[0m,\n",
       "        rating: \u001b[32m5.0\u001b[0m,\n",
       "        images: [\n",
       "          \u001b[32m\"https://s3.amazonaws.com/mikesskisport/images/Snow\"\u001b[0m\n",
       "        ]\n",
+      "      },\n",
+      "      {\n",
+      "        productId: \u001b[32m\"2\"\u001b[0m,\n",
+      "        name: \u001b[32m\"Mike's Ski Shop Exclusive\"\u001b[0m,\n",
+      "        description: \u001b[32m\"Mike's Ski Shop Exclusive\"\u001b[0m,\n",
+      "        category: \u001b[32m\"Ski Shops\"\u001b[0m,\n",
+      "        price: \u001b[32m20.0\u001b[0m,\n",
+      "        inStock: \u001b[32mTrue\u001b[0m,\n",
+      "        rating: \u001b[32m5.0\u001b[0m,\n",
+      "        images: [\n",
+      "          \u001b[32m\"https://s3.amazonaws.com/mikesskisport/images/Mike\"\u001b[0m\n",
+      "        ]\n",
+      "      },\n",
+      "      {\n",
+      "        productId: \u001b[32m\"3\"\u001b[0m,\n",
+      "        name: \u001b[32m\"Mike's Ski Shop Exclusive\"\u001b[0m,\n",
+      "        description: \u001b[32m\"Mike's Ski Shop Exclusive\"\u001b[0m,\n",
+      "        category: \u001b[32m\"Ski Shops\"\u001b[0m,\n",
+      "        price: \u001b[32m20.0\u001b[0m,\n",
+      "        inStock: \u001b[32mTrue\u001b[0m,\n",
+      "        rating: \u001b[32m5.0\u001b[0m,\n",
+      "        images: [\n",
+      "          \u001b[32m\"https://s3.amazonaws.com/mikesskisport/images/Mike\"\u001b[0m\n",
+      "        ]\n",
       "      }\n",
       "    ]\n",
       "  }\n",
@@ -175,62 +199,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Generating...\n",
-      "{\n",
-      "  make: \u001b[32m\"Mazda\"\u001b[0m,\n",
-      "  model: [\n",
-      "    {\n",
-      "      prob: \u001b[32m0.8154296875\u001b[0m,\n",
-      "      choice: \u001b[32m\"Kea\"\u001b[0m\n",
-      "    },\n",
-      "    {\n",
-      "      prob: \u001b[32m0.184814453125\u001b[0m,\n",
-      "      choice: \u001b[32m\"Mazda\"\u001b[0m\n",
-      "    }\n",
-      "  ],\n",
-      "  new: [\n",
-      "    {\n",
-      "      prob: \u001b[32m0.90185546875\u001b[0m,\n",
-      "      choice: \u001b[32m\"true\"\u001b[0m\n",
-      "    },\n",
-      "    {\n",
-      "      prob: \u001b[32m0.09808349609375\u001b[0m,\n",
-      "      choice: \u001b[32m\"false\"\u001b[0m\n",
-      "    }\n",
-      "  ],\n",
-      "  rating: [\n",
-      "    {\n",
-      "      prob: \u001b[32m0.221435546875\u001b[0m,\n",
-      "      choice: \u001b[32m\"1\"\u001b[0m\n",
-      "    },\n",
-      "    {\n",
-      "      prob: \u001b[32m0.394775390625\u001b[0m,\n",
-      "      choice: \u001b[32m\"2\"\u001b[0m\n",
-      "    },\n",
-      "    {\n",
-      "      prob: \u001b[32m0.382568359375\u001b[0m,\n",
-      "      choice: \u001b[32m\"3\"\u001b[0m\n",
-      "    },\n",
-      "    {\n",
-      "      prob: \u001b[32m0.0013370513916015625\u001b[0m,\n",
-      "      choice: \u001b[32m\"4\"\u001b[0m\n",
-      "    }\n",
-      "  ],\n",
-      "  year: \u001b[32m2016.0\u001b[0m,\n",
-      "  colors_available: [\n",
-      "    \u001b[32m\"red\"\u001b[0m\n",
-      "  ]\n",
-      "}\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "car = {\n",
     "    \"type\": \"object\",\n",
@@ -262,96 +233,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Generating...\n",
-      "{\n",
-      "  car: {\n",
-      "    make: \u001b[32m\"Rolls Royce\"\u001b[0m,\n",
-      "    model: \u001b[32m\"Phantom\"\u001b[0m,\n",
-      "    year: \u001b[32m2014.0\u001b[0m,\n",
-      "    colors: [\n",
-      "      {\n",
-      "        prob: \u001b[32m0.001560211181640625\u001b[0m,\n",
-      "        choice: \u001b[32m\"white\"\u001b[0m\n",
-      "      },\n",
-      "      {\n",
-      "        prob: \u001b[32m0.833984375\u001b[0m,\n",
-      "        choice: \u001b[32m\"red\"\u001b[0m\n",
-      "      },\n",
-      "      {\n",
-      "        prob: \u001b[32m0.0865478515625\u001b[0m,\n",
-      "        choice: \u001b[32m\"black\"\u001b[0m\n",
-      "      },\n",
-      "      {\n",
-      "        prob: \u001b[32m0.048553466796875\u001b[0m,\n",
-      "        choice: \u001b[32m\"blue\"\u001b[0m\n",
-      "      },\n",
-      "      {\n",
-      "        prob: \u001b[32m0.0294342041015625\u001b[0m,\n",
-      "        choice: \u001b[32m\"green\"\u001b[0m\n",
-      "      }\n",
-      "    ],\n",
-      "    as_new: [\n",
-      "      {\n",
-      "        prob: \u001b[32m0.96533203125\u001b[0m,\n",
-      "        choice: \u001b[32m\"true\"\u001b[0m\n",
-      "      },\n",
-      "      {\n",
-      "        prob: \u001b[32m0.03460693359375\u001b[0m,\n",
-      "        choice: \u001b[32m\"false\"\u001b[0m\n",
-      "      }\n",
-      "    ],\n",
-      "    rating: [\n",
-      "      {\n",
-      "        prob: \u001b[32m0.05462646484375\u001b[0m,\n",
-      "        choice: \u001b[32m\"1\"\u001b[0m\n",
-      "      },\n",
-      "      {\n",
-      "        prob: \u001b[32m0.233642578125\u001b[0m,\n",
-      "        choice: \u001b[32m\"2\"\u001b[0m\n",
-      "      },\n",
-      "      {\n",
-      "        prob: \u001b[32m0.2093505859375\u001b[0m,\n",
-      "        choice: \u001b[32m\"3\"\u001b[0m\n",
-      "      },\n",
-      "      {\n",
-      "        prob: \u001b[32m0.50244140625\u001b[0m,\n",
-      "        choice: \u001b[32m\"4\"\u001b[0m\n",
-      "      }\n",
-      "    ],\n",
-      "    features: {\n",
-      "      audio: {\n",
-      "        brand: \u001b[32m\"Mercedes-Benz\"\u001b[0m,\n",
-      "        speakers: \u001b[32m2.09999\u001b[0m,\n",
-      "        hasBluetooth: \u001b[32mTrue\u001b[0m\n",
-      "      },\n",
-      "      safety: {\n",
-      "        airbags: \u001b[32m2.09999\u001b[0m,\n",
-      "        parkingSensors: \u001b[32mTrue\u001b[0m,\n",
-      "        laneAssist: \u001b[32mTrue\u001b[0m\n",
-      "      },\n",
-      "      performance: {\n",
-      "        engine: \u001b[32m\"Mercedes-Benz 6.2 L\"\u001b[0m,\n",
-      "        horsepower: \u001b[32m423.09999\u001b[0m,\n",
-      "        topSpeed: \u001b[32m220.09999\u001b[0m\n",
-      "      }\n",
-      "    }\n",
-      "  },\n",
-      "  owner: {\n",
-      "    firstName: \u001b[32m\"John\"\u001b[0m,\n",
-      "    lastName: \u001b[32m\"Doe\"\u001b[0m,\n",
-      "    age: \u001b[32m40.09999\u001b[0m\n",
-      "  }\n",
-      "}\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "complex_car = {\n",
     "    \"type\": \"object\",\n",
@@ -418,6 +302,107 @@
     "\n",
     "highlight_values(output)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Readme example"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'name': 'John Doe',\n",
+       " 'age': 20,\n",
+       " 'age_probs': [{'prob': 0.794921875, 'choice': '12'},\n",
+       "  {'prob': 0.068359375, 'choice': '10'},\n",
+       "  {'prob': 0.04345703125, 'choice': '16'},\n",
+       "  {'prob': 0.03228759765625, 'choice': '14'},\n",
+       "  {'prob': 0.0175628662109375, 'choice': '11'},\n",
+       "  {'prob': 0.0157318115234375, 'choice': '15'},\n",
+       "  {'prob': 0.013671875, 'choice': '20'},\n",
+       "  {'prob': 0.006664276123046875, 'choice': '18'},\n",
+       "  {'prob': 0.0046539306640625, 'choice': '13'},\n",
+       "  {'prob': 0.0018215179443359375, 'choice': '21'},\n",
+       "  {'prob': 0.00041294097900390625, 'choice': '17'},\n",
+       "  {'prob': 0.00028824806213378906, 'choice': '19'},\n",
+       "  {'prob': 0.00014495849609375, 'choice': '22'},\n",
+       "  {'prob': 6.955862045288086e-05, 'choice': '23'},\n",
+       "  {'prob': 2.968311309814453e-05, 'choice': '25'},\n",
+       "  {'prob': 2.8789043426513672e-05, 'choice': '26'},\n",
+       "  {'prob': 1.901388168334961e-05, 'choice': '24'},\n",
+       "  {'prob': 1.1742115020751953e-05, 'choice': '28'},\n",
+       "  {'prob': 1.1920928955078125e-06, 'choice': '27'},\n",
+       "  {'prob': 7.748603820800781e-07, 'choice': '29'}],\n",
+       " 'unit_time': 0.01,\n",
+       " 'is_student': True,\n",
+       " 'is_student_probs': [{'prob': 0.8310546875, 'choice': 'true'},\n",
+       "  {'prob': 0.1688232421875, 'choice': 'false'}],\n",
+       " 'courses': ['C1'],\n",
+       " 'trim': None,\n",
+       " 'color': 'white'}"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from prob_jsonformer import Jsonformer\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "\n",
+    "model_name = \"databricks/dolly-v2-3b\"\n",
+    "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+    "\n",
+    "json_schema = {\n",
+    "    \"type\": \"object\",\n",
+    "    \"properties\": {\n",
+    "        \"name\": {\"type\": \"string\"},\n",
+    "        \"age\": {\"type\": \"integer\"},\n",
+    "        \"age_probs\": {\"type\": \"choice_probs\", \"enum\": [str(s) for s in range(10, 30)]},\n",
+    "        \"unit_time\": {\"type\": \"number\"},\n",
+    "        \"is_student\": {\"type\": \"boolean\"},\n",
+    "        \"is_student_probs\": {\"type\": \"choice_probs\", \"enum\": [\"true\", \"false\"]},\n",
+    "        \"courses\": {\n",
+    "            \"type\": \"array\",\n",
+    "            \"items\": {\"type\": \"string\"}\n",
+    "        },\n",
+    "        \"trim\": {\"type\": [\"string\", \"null\"]},\n",
+    "        \"color\": {\n",
+    "            \"type\": \"enum\",\n",
+    "            \"values\": [\"red\", \"green\", \"blue\", \"brown\", \"white\", \"black\"],\n",
+    "        },\n",
+    "    }\n",
+    "}\n",
+    "\n",
+    "prompt = \"Generate a young person's information based on the following schema:\"\n",
+    "jsonformer = Jsonformer(model, tokenizer, json_schema, prompt, temperature=0)\n",
+    "generated_data = jsonformer()\n",
+    "\n",
+    "generated_data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
diff --git a/prob_jsonformer/choice_tree.py b/prob_jsonformer/choice_tree.py
index 61a731d..83ee70a 100644
--- a/prob_jsonformer/choice_tree.py
+++ b/prob_jsonformer/choice_tree.py
@@ -21,7 +21,7 @@ def get_valid_next_choices(choices_tokens, current_tokens):
     return torch.LongTensor(next_choices)
 
 
-def choice_tree(
+def _choice_tree(
     model: AutoModelForCausalLM,
     tokenizer: AutoTokenizer,
     input_ids: Int[Tensor, "seq"],
@@ -29,7 +29,6 @@ def choice_tree(
     choice: Optional[Int[Tensor, ""]] = None,
     prob: float = 1,
     current_tokens: Int[Tensor, "seq"] = torch.LongTensor([]),
-    z=[],
 ):
     if choice is not None:
         c = choice[None].to(current_tokens.device)
@@ -57,5 +56,19 @@ def choice_tree(
                 choice=next_choice,
                 prob=next_prob,
                 current_tokens=current_tokens,
-                z=z + [i],
             )
+
+
+def choice_tree(
+    *args,
+    **kwargs,
+):
+    choice_json = list(
+        _choice_tree(
+            *args,
+            **kwargs,
+        )
+    )
+    # order by probability
+    choice_json = sorted(choice_json, key=lambda x: -x["prob"])
+    return choice_json
diff --git a/prob_jsonformer/main.py b/prob_jsonformer/main.py
index dd386a2..d23b0ab 100644
--- a/prob_jsonformer/main.py
+++ b/prob_jsonformer/main.py
@@ -138,10 +138,7 @@ class Jsonformer:
         true_token_id = self.tokenizer.encode("true", return_tensors="pt")[0, 0]
         false_token_id = self.tokenizer.encode("false", return_tensors="pt")[0, 0]
 
-        if prob:
-            result = dict(true=logits[true_token_id], false=logits[false_token_id])
-        else:
-            result = logits[true_token_id] > logits[false_token_id]
+        result = logits[true_token_id] > logits[false_token_id]
 
         self.debug("[generate_boolean]", result)
 
@@ -201,24 +198,38 @@ class Jsonformer:
         self.debug("[generate_enum]", prompt, is_prompt=True)
 
         # These are necessary because we don't know if we're at the end or middle of an object/array
-        terminal_tokens = torch.concat([
-            self.tokenizer.encode(s, add_special_tokens=False, return_tensors="pt")[:, 0]
-            for s in ('", "', '"}', '"]')
-        ])
+        terminal_tokens = torch.concat(
+            [
+                self.tokenizer.encode(s, add_special_tokens=False, return_tensors="pt")[
+                    :, 0
+                ]
+                for s in ('", "', '"}', '"]')
+            ]
+        )
 
         highest_probability = 0.0
         best_option = None
         for option in enum_values:
-            n_option_tokens = self.tokenizer.encode(f'"{option}', add_special_tokens=False, return_tensors="pt").shape[1]
-            prompt_tokens = self.tokenizer.encode(prompt + f'"{option}', return_tensors="pt")
+            n_option_tokens = self.tokenizer.encode(
+                f'"{option}', add_special_tokens=False, return_tensors="pt"
+            ).shape[1]
+            prompt_tokens = self.tokenizer.encode(
+                prompt + f'"{option}', return_tensors="pt"
+            )
             option_tokens = prompt_tokens[0, -n_option_tokens:]
 
             with torch.no_grad():
-                logits = self.model.forward(prompt_tokens.to(self.model.device)).logits[0, -n_option_tokens-1:]
+                logits = self.model.forward(prompt_tokens.to(self.model.device)).logits[
+                    0, -n_option_tokens - 1 :
+                ]
             probabilities = torch.softmax(logits, dim=1)
-            option_token_probabilities = probabilities[:-1][torch.arange(probabilities.shape[0]-1), option_tokens]
+            option_token_probabilities = probabilities[:-1][
+                torch.arange(probabilities.shape[0] - 1), option_tokens
+            ]
             termination_probability = torch.max(probabilities[-1, terminal_tokens])
-            option_probability = torch.prod(option_token_probabilities) * termination_probability
+            option_probability = (
+                torch.prod(option_token_probabilities) * termination_probability
+            )
 
             if option_probability > highest_probability:
                 best_option = option
@@ -252,7 +263,7 @@ class Jsonformer:
         max_type = None
         max_logit = -float("inf")
         for possible_type in possible_types:
-            try: 
+            try:
                 prefix_tokens = self.type_prefix_tokens[possible_type]
             except KeyError:
                 raise ValueError(f"Unsupported schema type: {possible_type}")
@@ -265,7 +276,7 @@ class Jsonformer:
             raise Exception("Unable to find best type to generate for union type")
         self.debug("[choose_type_to_generate]", max_type)
         return max_type
-    
+
     def generate_value(
         self,
         schema: Dict[str, Any],