{
    "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": "stderr",
                    "output_type": "stream",
                    "text": [
                        "/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/.venv/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"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Loading model and tokenizer...\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "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",
                        "/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"
                    ]
                }
            ],
            "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": null,
            "metadata": {},
            "outputs": [],
            "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": 9,
            "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": ".venv",
            "language": "python",
            "name": "python3"
        },
        "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"
        },
        "orig_nbformat": 4
    },
    "nbformat": 4,
    "nbformat_minor": 2
}