Spaces:
Paused
Paused
File size: 4,500 Bytes
0538b40 5f146f4 0538b40 42bc72f 0538b40 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
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()
|