{
 "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"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c7620cdf45c54ab5abaf479b141f4479",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/819 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b9dba56663274820997ef5a07e7acea4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/5.68G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8571f2926f9541469ed34140333cdf81",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/5.68G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "93c30a28fbdd4257ab783cc659b24687",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/450 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "10ad2621da06464bbdc5b4aa49687d50",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/2.11M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3abe2819228e4da4ad859c54abefc479",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/228 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "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 29\u001b[0m\n\u001b[1;32m     27\u001b[0m prompt \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mGenerate a young person\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124ms information based on the following schema:\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     28\u001b[0m jsonformer \u001b[38;5;241m=\u001b[39m Jsonformer(model, tokenizer, json_schema, prompt)\n\u001b[0;32m---> 29\u001b[0m generated_data \u001b[38;5;241m=\u001b[39m \u001b[43mjsonformer\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     31\u001b[0m generated_data\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: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 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\", \"enum\": [str(s) for s in range(10, 30)]},\n",
    "        # we can return the probabilistic weighted mean of a range\n",
    "        \"age_wmean\": {\"type\": \"p_integer\", \"minimum\": 10, \"maximum\": 30},\n",
    "        # the prob of true and false\n",
    "        \"is_student_probs\": {\"type\": \"p_enum\", \"enum\": [\"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",
    "\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": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_data = {\n",
    "    \"age_probs\": [\n",
    "        {\"prob\": 0.94091796875, \"choice\": \"10\"},\n",
    "        {\"prob\": 0.033233642578125, \"choice\": \"20\"},\n",
    "        {\"prob\": 0.0122222900390625, \"choice\": \"12\"},\n",
    "        {\"prob\": 0.00412750244140625, \"choice\": \"21\"},\n",
    "        {\"prob\": 0.0028362274169921875, \"choice\": \"16\"},\n",
    "        {\"prob\": 0.0018453598022460938, \"choice\": \"15\"},\n",
    "        {\"prob\": 0.00113677978515625, \"choice\": \"11\"},\n",
    "        {\"prob\": 0.0011110305786132812, \"choice\": \"18\"},\n",
    "        {\"prob\": 0.0005083084106445312, \"choice\": \"25\"},\n",
    "        {\"prob\": 0.0004558563232421875, \"choice\": \"23\"},\n",
    "        {\"prob\": 0.0002498626708984375, \"choice\": \"14\"},\n",
    "        {\"prob\": 0.00023281574249267578, \"choice\": \"13\"},\n",
    "        {\"prob\": 0.0002238750457763672, \"choice\": \"22\"},\n",
    "        {\"prob\": 0.00018131732940673828, \"choice\": \"26\"},\n",
    "        {\"prob\": 0.0001690387725830078, \"choice\": \"24\"},\n",
    "        {\"prob\": 0.00012552738189697266, \"choice\": \"19\"},\n",
    "        {\"prob\": 7.796287536621094e-05, \"choice\": \"27\"},\n",
    "        {\"prob\": 7.265806198120117e-05, \"choice\": \"28\"},\n",
    "        {\"prob\": 4.106760025024414e-05, \"choice\": \"17\"},\n",
    "        {\"prob\": 2.5033950805664062e-06, \"choice\": \"29\"},\n",
    "    ],\n",
    "    \"age_wmean\": 17.816404402256012,\n",
    "    \"is_student_probs\": [\n",
    "        {\"prob\": 0.974609375, \"choice\": \"true\"},\n",
    "        {\"prob\": 0.025177001953125, \"choice\": \"false\"},\n",
    "    ],\n",
    "    \"is_student\": False,\n",
    "    \"name\": \"John\",\n",
    "    \"age\": 17,\n",
    "    \"unit_time\": 0.5,\n",
    "    \"courses\": [\"C++\"],\n",
    "    \"trim\": None,\n",
    "    \"color\": \"white\",\n",
    "}"
   ]
  },
  {
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}