import json
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from jsonformer.format import highlight_values
from jsonformer.main import Jsonformer

print("Loading model and tokenizer...")
model_name = "databricks/dolly-v2-3b"
model = AutoModelForCausalLM.from_pretrained(model_name, use_cache=True, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, use_cache=True)
print("Loaded model and tokenizer")

def generate(input_prompt, input_schema):
    try:
        if not input_prompt:
            raise ValueError("Prompt is empty")
        if not input_schema:
            raise ValueError("JSON Schema is empty")
        input_schema = json.loads(input_schema)
        builder = Jsonformer(
            model=model,
            tokenizer=tokenizer,
            json_schema=input_schema,
            prompt=input_prompt,
        )
        print("Generating...")
        output_json = builder()
        return output_json
    except Exception as e:
        raise gr.Error(e)

examples = [
    [
        "Generate a json where it is silver Aston Martin DB5 manufactured in 1964",
        '{\n  "type": "object",\n  "properties": {\n    "car": {\n      "type": "object",\n      "properties": {\n        "make": {\n          "type": "string"\n        },\n        "model": {\n          "type": "string"\n        },\n        "year": {\n          "type": "number"\n        },\n        "colors": {\n          "type": "array",\n          "items": {\n            "type": "string"\n          }\n        }\n      }\n    }\n  }\n}'
    ],
    [
        "Generate a person's information based on the following schema. The person is Lionel Messi, aged 26. Messi is a student at Georgia Tech, and take the following courses: Chemistry, Mathematics, and a minor in Japanese.",
        '{\n  "type": "object",\n  "properties": {\n    "name": {\n      "type": "string"\n    },\n    "age": {\n      "type": "number"\n    },\n    "is_student": {\n      "type": "boolean"\n    },\n    "courses": {\n      "type": "array",\n      "items": {\n        "type": "string"\n      }\n    }\n  }\n}'
    ],
]

css = """
#examples {
    width: 35rem;
}
"""

with gr.Blocks(css=css) as demo:
    gr.HTML(
        """
            <div style="text-align: center; margin: 0 auto;">
              <div
                style="
                  display: inline-flex;
                  align-items: center;
                  gap: 0.8rem;
                  font-size: 1.75rem;
                "
              >
                <h1 style="font-weight: 900; margin-bottom: 7px;margin-top:5px">
                  Jsonformer
                </h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">
                <a
                  href="https://github.com/1rgs/jsonformer"
                  style="text-decoration: underline;"
                  target="_blank"
                  >Jsonformer</a>: A Bulletproof Way to Generate Structured JSON from Language Models.
              </p>
              <p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">
                Jsonformer generates <b>syntactically correct</b> jsons by constraining/shrinking output space of Language Models.
              </p>
            </div>
        """
    )
    with gr.Row():
        with gr.Column(scale=1, min_width=600):
            input_prompt = gr.TextArea("Generate a json where it is silver Aston Martin DB5 manufactured in 1964", label="Prompt", lines=2)
            input_schema = gr.Code('{\n  "type": "object",\n  "properties": {\n    "car": {\n      "type": "object",\n      "properties": {\n        "make": {\n          "type": "string"\n        },\n        "model": {\n          "type": "string"\n        },\n        "year": {\n          "type": "number"\n        },\n        "colors": {\n          "type": "array",\n          "items": {\n            "type": "string"\n          }\n        }\n      }\n    }\n  }\n}', label="JSON Schema")
            generate_btn = gr.Button("Generate")
        with gr.Column(scale=1, min_width=600):
            output_json = gr.JSON(label="Generated JSON")
    
    ex = gr.Examples(examples=examples, fn=generate, inputs=[input_prompt, input_schema], outputs=output_json, cache_examples=False, elem_id="examples",)
    ex.dataset.headers = [""]
    generate_btn.click(fn=generate, inputs=[input_prompt, input_schema], outputs=output_json, api_name="greet")

demo.launch()