diff --git a/example.ipynb b/example.ipynb
index 825b797..49520e8 100644
--- a/example.ipynb
+++ b/example.ipynb
@@ -37,37 +37,42 @@
    "execution_count": 2,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Generating...\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
-     ]
-    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "Generating...\n",
       "{\n",
-      "  name: \u001b[32m\"Mickey Mouse\"\u001b[0m,\n",
-      "  age: \u001b[32m30.0\u001b[0m,\n",
-      "  description: \u001b[32m\"A mouse like object\"\u001b[0m\n",
+      "  store: {\n",
+      "    name: \u001b[32m\"Mike's Ski Shop\"\u001b[0m,\n",
+      "    location: \u001b[32m\"Sugarloaf\"\u001b[0m,\n",
+      "    inventory: [\n",
+      "      {\n",
+      "        productId: \u001b[32m\"1234567890\"\u001b[0m,\n",
+      "        name: \u001b[32m\"Ski Shop\"\u001b[0m,\n",
+      "        description: \u001b[32m\"Ski Shop sells premium skis and snowboards\"\u001b[0m,\n",
+      "        category: \u001b[32m\"Sports\"\u001b[0m,\n",
+      "        price: \u001b[32m20.09999\u001b[0m,\n",
+      "        inStock: \u001b[32mTrue\u001b[0m,\n",
+      "        rating: \u001b[32m5.09999\u001b[0m,\n",
+      "        images: [\n",
+      "          \u001b[32m\"https://s3.amazonaws.com/images.skisnow.com/skis\"\u001b[0m\n",
+      "        ]\n",
+      "      },\n",
+      "      {\n",
+      "        productId: \u001b[32m\"12345678910\"\u001b[0m,\n",
+      "        name: \u001b[32m\"Snowboard\"\u001b[0m,\n",
+      "        description: \u001b[32m\"Snowboard for sale\"\u001b[0m,\n",
+      "        category: \u001b[32m\"Sports\"\u001b[0m,\n",
+      "        price: \u001b[32m20.09999\u001b[0m,\n",
+      "        inStock: \u001b[32mTrue\u001b[0m,\n",
+      "        rating: \u001b[32m5.09999\u001b[0m,\n",
+      "        images: [\n",
+      "          \u001b[32m\"https://s3.amazonaws.com/images.skisnow.com/snow\"\u001b[0m\n",
+      "        ]\n",
+      "      }\n",
+      "    ]\n",
+      "  }\n",
       "}\n"
      ]
     }
@@ -76,27 +81,48 @@
     "from jsonformer.format import highlight_values\n",
     "from jsonformer.main import Jsonformer\n",
     "\n",
-    "mickey_mouse_schema = {\n",
+    "ecomm = {\n",
     "    \"type\": \"object\",\n",
     "    \"properties\": {\n",
-    "        \"name\": {\"type\": \"string\"},\n",
-    "        \"age\": {\"type\": \"number\"},\n",
-    "        \"description\": {\"type\": \"string\"},\n",
+    "        \"store\": {\n",
+    "            \"type\": \"object\",\n",
+    "            \"properties\": {\n",
+    "                \"name\": {\"type\": \"string\"},\n",
+    "                \"location\": {\"type\": \"string\"},\n",
+    "                \"inventory\": {\n",
+    "                    \"type\": \"array\",\n",
+    "                    \"items\": {\n",
+    "                        \"type\": \"object\",\n",
+    "                        \"properties\": {\n",
+    "                            \"productId\": {\"type\": \"string\"},\n",
+    "                            \"name\": {\"type\": \"string\"},\n",
+    "                            \"description\": {\"type\": \"string\"},\n",
+    "                            \"category\": {\"type\": \"string\"},\n",
+    "                            \"price\": {\"type\": \"number\"},\n",
+    "                            \"inStock\": {\"type\": \"boolean\"},\n",
+    "                            \"rating\": {\"type\": \"number\"},\n",
+    "                            \"images\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}},\n",
+    "                        },\n",
+    "                    },\n",
+    "                },\n",
+    "            },\n",
+    "        }\n",
     "    },\n",
     "}\n",
     "\n",
+    "\n",
     "builder = Jsonformer(\n",
     "    model=model,\n",
     "    tokenizer=tokenizer,\n",
-    "    json_schema=mickey_mouse_schema,\n",
-    "    prompt=\"write a description about mickey mouse\",\n",
-    "    max_string_token_length=100,\n",
+    "    json_schema=ecomm,\n",
+    "    prompt=\"write a description about mike's ski shop which sells premium skis and snowboards\",\n",
+    "    max_string_token_length=20,\n",
     ")\n",
     "\n",
     "print(\"Generating...\")\n",
     "output = builder()\n",
     "\n",
-    "highlight_values(output)\n"
+    "highlight_values(output)"
    ]
   },
   {
@@ -104,35 +130,11 @@
    "execution_count": 3,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Generating...\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
-     ]
-    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "Generating...\n",
       "{\n",
       "  make: \u001b[32m\"audi\"\u001b[0m,\n",
       "  model: \u001b[32m\"a4\"\u001b[0m,\n",
@@ -176,43 +178,11 @@
    "execution_count": 4,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Generating...\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
-      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
-     ]
-    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "Generating...\n",
       "{\n",
       "  car: {\n",
       "    make: \u001b[32m\"Rolls Royce\"\u001b[0m,\n",
diff --git a/jsonformer/main.py b/jsonformer/main.py
index 62b2b09..dd867d4 100644
--- a/jsonformer/main.py
+++ b/jsonformer/main.py
@@ -135,9 +135,7 @@ class Jsonformer:
         self.debug("[generate_string]", "|" + response + "|")
 
         if response.count('"') < 1:
-            raise ValueError(
-                "Failed to generate a valid string, try increasing max_string_token_length"
-            )
+            return response
 
         return response.split('"')[0].strip()
 
@@ -201,12 +199,23 @@ class Jsonformer:
             output = self.model.forward(input_tensor.to(self.model.device))
             logits = output.logits[0, -1]
 
-            close_bracket_token_id = self.tokenizer.convert_tokens_to_ids("]")
-            comma_token_id = self.tokenizer.convert_tokens_to_ids(", ")
-            close_bracket_logits = logits[close_bracket_token_id]
-            comma_logits = logits[comma_token_id]
 
-            if close_bracket_logits > comma_logits:
+            top_indices = logits.topk(30).indices
+            sorted_token_ids = top_indices[logits[top_indices].argsort(descending=True)]
+
+            found_comma = False
+            found_close_bracket = False
+
+            for token_id in sorted_token_ids:
+                decoded_token = self.tokenizer.decode(token_id)
+                if ',' in decoded_token:
+                    found_comma = True
+                    break
+                if ']' in decoded_token:
+                    found_close_bracket = True
+                    break
+
+            if found_close_bracket or not found_comma:
                 break
 
         return obj