Create app.py

app.py

@@ -0,0 +1,90 @@
+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. Mishig 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(label="Prompt", lines=2)
+            input_schema = gr.Code(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()