Jsonformer_example.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/1rgs/jsonformer/blob/main/Jsonformer_example.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "p-kogdhHO3PC",
        "outputId": "d41737a8-998e-4451-cce3-6ac6c98cc69b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.28.1)\n",
            "Requirement already satisfied: accelerate in /usr/local/lib/python3.10/dist-packages (0.18.0)\n",
            "Requirement already satisfied: jsonformer in /usr/local/lib/python3.10/dist-packages (0.9.0)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.65.0)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.14.1)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.27.1)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (23.1)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.22.4)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2022.10.31)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.12.0)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.13.3)\n",
            "Requirement already satisfied: torch>=1.4.0 in /usr/local/lib/python3.10/dist-packages (from accelerate) (2.0.0+cu118)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from accelerate) (5.9.5)\n",
            "Requirement already satisfied: termcolor<3.0.0,>=2.3.0 in /usr/local/lib/python3.10/dist-packages (from jsonformer) (2.3.0)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.11.0->transformers) (2023.4.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.11.0->transformers) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.4.0->accelerate) (1.11.1)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.4.0->accelerate) (3.1)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.4.0->accelerate) (2.0.0)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.4.0->accelerate) (3.1.2)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.4.0->accelerate) (16.0.2)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.4.0->accelerate) (3.25.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.4)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (1.26.15)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.0.12)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2022.12.7)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.4.0->accelerate) (2.1.2)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.4.0->accelerate) (1.3.0)\n"
          ]
        }
      ],
      "source": [
        "!pip install transformers accelerate jsonformer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9y0Jjj0sOtor",
        "outputId": "6c5ea465-384f-4aa1-8304-6fa721534f73"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading model and tokenizer...\n",
            "Loaded model and tokenizer\n"
          ]
        }
      ],
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "\n",
        "print(\"Loading model and tokenizer...\")\n",
        "model_name = \"databricks/dolly-v2-3b\"\n",
        "model = AutoModelForCausalLM.from_pretrained(model_name, use_cache=True, device_map=\"auto\")\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, use_cache=True)\n",
        "print(\"Loaded model and tokenizer\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8E7zFQ3NO0Yd",
        "outputId": "5f4d1ad0-1acf-48e0-d9bd-82f88ebcf97a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generating...\n",
            "{\n",
            "  car: {\n",
            "    make: \"audi\",\n",
            "    model: \"model A8\",\n",
            "    year: 2016.0,\n",
            "    colors: [\n",
            "      \"blue\"\n",
            "    ],\n",
            "    features: {\n",
            "      audio: {\n",
            "        brand: \"sony\",\n",
            "        speakers: 2.0,\n",
            "        hasBluetooth: True\n",
            "      },\n",
            "      safety: {\n",
            "        airbags: 2.0,\n",
            "        parkingSensors: True,\n",
            "        laneAssist: True\n",
            "      },\n",
            "      performance: {\n",
            "        engine: \"4.0\",\n",
            "        horsepower: 220.0,\n",
            "        topSpeed: 220.0\n",
            "      }\n",
            "    }\n",
            "  },\n",
            "  owner: {\n",
            "    firstName: \"John\",\n",
            "    lastName: \"Doe\",\n",
            "    age: 40.0\n",
            "  }\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "from jsonformer.format import highlight_values\n",
        "from jsonformer.main import Jsonformer\n",
        "\n",
        "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\": \"array\",\n",
        "          \"items\": {\"type\": \"string\"}\n",
        "        },\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",
        "\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)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "0eA8hzFHA1jo"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "premium",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}