{
    "cells": [
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {},
            "outputs": [],
            "source": [
                "# autoreload your package\n",
                "%load_ext autoreload\n",
                "%autoreload 2"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Loading model and tokenizer...\n",
                        "Loaded model and tokenizer\n"
                    ]
                }
            ],
            "source": [
                "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
                "import torch\n",
                "\n",
                "print(\"Loading model and tokenizer...\")\n",
                "model_name = \"databricks/dolly-v2-3b\"\n",
                "model = AutoModelForCausalLM.from_pretrained(\n",
                "    model_name,\n",
                "    use_cache=True,\n",
                "    torch_dtype=torch.float16,\n",
                "    attn_implementation=\"eager\",\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": [
                "# Continue"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Generating...\n"
                    ]
                },
                {
                    "ename": "KeyError",
                    "evalue": "'values'",
                    "output_type": "error",
                    "traceback": [
                        "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
                        "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
                        "Cell \u001b[0;32mIn[3], line 47\u001b[0m\n\u001b[1;32m     38\u001b[0m builder \u001b[38;5;241m=\u001b[39m Jsonformer(\n\u001b[1;32m     39\u001b[0m     model\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[1;32m     40\u001b[0m     tokenizer\u001b[38;5;241m=\u001b[39mtokenizer,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     43\u001b[0m     max_string_token_length\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m20\u001b[39m,\n\u001b[1;32m     44\u001b[0m )\n\u001b[1;32m     46\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mGenerating...\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 47\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[43mbuilder\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     49\u001b[0m highlight_values(output)\n",
                        "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:439\u001b[0m, in \u001b[0;36mJsonformer.__call__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    437\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Dict[\u001b[38;5;28mstr\u001b[39m, Any]:\n\u001b[1;32m    438\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvalue \u001b[38;5;241m=\u001b[39m {}\n\u001b[0;32m--> 439\u001b[0m     generated_data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgenerate_object\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    440\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mjson_schema\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mproperties\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m\n\u001b[1;32m    441\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    442\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m generated_data\n",
                        "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:274\u001b[0m, in \u001b[0;36mJsonformer.generate_object\u001b[0;34m(self, properties, obj)\u001b[0m\n\u001b[1;32m    272\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, schema \u001b[38;5;129;01min\u001b[39;00m properties\u001b[38;5;241m.\u001b[39mitems():\n\u001b[1;32m    273\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m[generate_object] generating value for\u001b[39m\u001b[38;5;124m\"\u001b[39m, key)\n\u001b[0;32m--> 274\u001b[0m     obj[key] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgenerate_value\u001b[49m\u001b[43m(\u001b[49m\u001b[43mschema\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mobj\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    275\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m obj\n",
                        "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:376\u001b[0m, in \u001b[0;36mJsonformer.generate_value\u001b[0;34m(self, schema, obj, key)\u001b[0m\n\u001b[1;32m    374\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    375\u001b[0m         obj\u001b[38;5;241m.\u001b[39mappend(new_obj)\n\u001b[0;32m--> 376\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgenerate_object\u001b[49m\u001b[43m(\u001b[49m\u001b[43mschema\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mproperties\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnew_obj\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    377\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m schema_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnull\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m    378\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
                        "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:274\u001b[0m, in \u001b[0;36mJsonformer.generate_object\u001b[0;34m(self, properties, obj)\u001b[0m\n\u001b[1;32m    272\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, schema \u001b[38;5;129;01min\u001b[39;00m properties\u001b[38;5;241m.\u001b[39mitems():\n\u001b[1;32m    273\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m[generate_object] generating value for\u001b[39m\u001b[38;5;124m\"\u001b[39m, key)\n\u001b[0;32m--> 274\u001b[0m     obj[key] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgenerate_value\u001b[49m\u001b[43m(\u001b[49m\u001b[43mschema\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mobj\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    275\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m obj\n",
                        "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:351\u001b[0m, in \u001b[0;36mJsonformer.generate_value\u001b[0;34m(self, schema, obj, key)\u001b[0m\n\u001b[1;32m    349\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    350\u001b[0m         obj\u001b[38;5;241m.\u001b[39mappend(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgeneration_marker)\n\u001b[0;32m--> 351\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgenerate_p_enum(\u001b[43mschema\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mvalues\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m, \u001b[38;5;28mround\u001b[39m\u001b[38;5;241m=\u001b[39mschema\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mround\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;241m3\u001b[39m))\n\u001b[1;32m    352\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m schema_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mp_integer\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m    353\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m key:\n",
                        "\u001b[0;31mKeyError\u001b[0m: 'values'"
                    ]
                }
            ],
            "source": [
                "from prob_jsonformer.format import highlight_values\n",
                "from prob_jsonformer.main import Jsonformer\n",
                "\n",
                "ecomm = {\n",
                "    \"type\": \"object\",\n",
                "    \"properties\": {\n",
                "        \"store\": {\n",
                "            \"type\": \"object\",\n",
                "            \"properties\": {\n",
                "                \"name\": {\"type\": \"string\"},\n",
                "                \"location\": {\"type\": \"string\"},\n",
                "                \"p_enum\": {\n",
                "                    \"type\": \"p_enum\",\n",
                "                    \"enum\": [\"ski\", \"snowboard\", \"walk\", \"pretend\"],\n",
                "                },\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=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)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "car = {\n",
                "    \"type\": \"object\",\n",
                "    \"properties\": {\n",
                "        \"make\": {\"type\": \"string\"},\n",
                "        \"model\": {\"type\": \"p_enum\", \"enum\": [\"Mazda\", \"Kea\"]},\n",
                "        \"new\": {\"type\": \"p_enum\", \"enum\": [\"true\", \"false\"]},\n",
                "        \"rating\": {\"type\": \"p_enum\", \"enum\": [\"1\", \"2\", \"3\", \"4\"]},\n",
                "        \"year\": {\"type\": \"number\"},\n",
                "        \"colors_available\": {\n",
                "            \"type\": \"array\",\n",
                "            \"items\": {\"type\": \"string\"},\n",
                "        },\n",
                "    },\n",
                "}\n",
                "\n",
                "builder = Jsonformer(\n",
                "    model=model,\n",
                "    tokenizer=tokenizer,\n",
                "    json_schema=car,\n",
                "    prompt=\"generate an example car\",\n",
                ")\n",
                "\n",
                "print(\"Generating...\")\n",
                "output = builder()\n",
                "\n",
                "highlight_values(output)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "complex_car = {\n",
                "    \"type\": \"object\",\n",
                "    \"properties\": {\n",
                "        \"car\": {\n",
                "            \"type\": \"object\",\n",
                "            \"properties\": {\n",
                "                \"make\": {\"type\": \"string\"},\n",
                "                \"model\": {\"type\": \"string\"},\n",
                "                \"year\": {\"type\": \"number\"},\n",
                "                \"colors\": {\n",
                "                    \"type\": \"p_enum\",\n",
                "                    \"enum\": [\"red\", \"green\", \"blue\", \"black\", \"white\"],\n",
                "                },\n",
                "                \"as_new\": {\"type\": \"p_enum\", \"enum\": [\"true\", \"false\"]},\n",
                "                \"rating\": {\"type\": \"p_enum\", \"enum\": [\"1\", \"2\", \"3\", \"4\"]},\n",
                "                \"features\": {\n",
                "                    \"type\": \"object\",\n",
                "                    \"properties\": {\n",
                "                        \"audio\": {\n",
                "                            \"type\": \"object\",\n",
                "                            \"properties\": {\n",
                "                                \"brand\": {\"type\": \"string\"},\n",
                "                                \"speakers\": {\"type\": \"number\"},\n",
                "                                \"hasBluetooth\": {\"type\": \"boolean\"},\n",
                "                            },\n",
                "                        },\n",
                "                        \"safety\": {\n",
                "                            \"type\": \"object\",\n",
                "                            \"properties\": {\n",
                "                                \"airbags\": {\"type\": \"number\"},\n",
                "                                \"parkingSensors\": {\"type\": \"boolean\"},\n",
                "                                \"laneAssist\": {\"type\": \"boolean\"},\n",
                "                            },\n",
                "                        },\n",
                "                        \"performance\": {\n",
                "                            \"type\": \"object\",\n",
                "                            \"properties\": {\n",
                "                                \"engine\": {\"type\": \"string\"},\n",
                "                                \"horsepower\": {\"type\": \"number\"},\n",
                "                                \"topSpeed\": {\"type\": \"number\"},\n",
                "                            },\n",
                "                        },\n",
                "                    },\n",
                "                },\n",
                "            },\n",
                "        },\n",
                "        \"owner\": {\n",
                "            \"type\": \"object\",\n",
                "            \"properties\": {\n",
                "                \"firstName\": {\"type\": \"string\"},\n",
                "                \"lastName\": {\"type\": \"string\"},\n",
                "                \"age\": {\"type\": \"number\"},\n",
                "            },\n",
                "        },\n",
                "    },\n",
                "}\n",
                "builder = Jsonformer(\n",
                "    model=model,\n",
                "    tokenizer=tokenizer,\n",
                "    json_schema=complex_car,\n",
                "    prompt=\"generate an example Rolls Royce Phantom\",\n",
                ")\n",
                "\n",
                "print(\"Generating...\")\n",
                "output = builder()\n",
                "\n",
                "highlight_values(output)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Readme example"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "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)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "{'age_probs': [{'prob': 0.62353515625, 'choice': '10'},\n",
                            "  {'prob': 0.349609375, 'choice': '12'},\n",
                            "  {'prob': 0.01123809814453125, 'choice': '11'},\n",
                            "  {'prob': 0.00760650634765625, 'choice': '16'},\n",
                            "  {'prob': 0.0025482177734375, 'choice': '13'},\n",
                            "  {'prob': 0.0025081634521484375, 'choice': '15'},\n",
                            "  {'prob': 0.0018062591552734375, 'choice': '14'},\n",
                            "  {'prob': 0.00104522705078125, 'choice': '18'},\n",
                            "  {'prob': 0.00011551380157470703, 'choice': '17'},\n",
                            "  {'prob': 5.042552947998047e-05, 'choice': '19'}],\n",
                            " 'age_wmean': 15.544570922851562,\n",
                            " 'is_student_probs': [{'prob': 0.962890625, 'choice': 'true'},\n",
                            "  {'prob': 0.037322998046875, 'choice': 'false'}],\n",
                            " 'is_student': False,\n",
                            " 'name': 'John',\n",
                            " 'age': 17,\n",
                            " 'unit_time': 0.5,\n",
                            " 'courses': ['C++'],\n",
                            " 'trim': None,\n",
                            " 'color': 'green'}"
                        ]
                    },
                    "execution_count": 9,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "from prob_jsonformer import Jsonformer\n",
                "\n",
                "json_schema = {\n",
                "    \"type\": \"object\",\n",
                "    \"properties\": {\n",
                "        # we can return the probability of each choice, even if they are multiple tokens\n",
                "        \"age_probs\": {\"type\": \"p_enum\", \"values\": [str(s) for s in range(10, 20)]},\n",
                "        # we can return the probabilistic weighted mean of a range\n",
                "        \"age_wmean\": {\"type\": \"p_integer\", \"minimum\": 10, \"maximum\": 20},\n",
                "        # the prob of true and false\n",
                "        \"is_student_probs\": {\"type\": \"p_enum\", \"values\": [\"true\", \"false\"]},\n",
                "        \"is_student\": {\"type\": \"boolean\"},\n",
                "        # we've merged patches for enum, integer, null, union - currently mising from jsonformer\n",
                "        \"name\": {\"type\": \"string\", \"maxLength\": 4},\n",
                "        \"age\": {\"type\": \"integer\"},\n",
                "        \"unit_time\": {\"type\": \"number\"},\n",
                "        \"courses\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}},\n",
                "        \"trim\": {\"type\": [\"string\", \"null\"]},\n",
                "        \"color\": {\n",
                "            \"type\": \"enum\",\n",
                "            \"values\": [\"red\", \"green\", \"blue\", \"brown\", \"white\", \"black\"],\n",
                "        },\n",
                "    },\n",
                "}\n",
                "prompt = \"Generate a young person's information based on the following schema:\"\n",
                "jsonformer = Jsonformer(model, tokenizer, json_schema, prompt)\n",
                "generated_data = jsonformer()\n",
                "\n",
                "generated_data"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": []
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "pytorch_hf_env",
            "language": "python",
            "name": "pytorch_hf_env"
        },
        "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.10.16"
        },
        "orig_nbformat": 4
    },
    "nbformat": 4,
    "nbformat_minor": 2
}