Spaces:
Paused
Paused
| 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() | |