Ferr0/structured-output-playground

1"""Structured Output Playground — lock any LLM's output to a JSON schema.2 3A local model (Qwen2.5-3B-Instruct) extracts structured data from free text. With4*constrained decoding* on (Outlines), the decoder can only emit tokens that keep the5output conformant to the schema — right keys, right types, valid enums, every time.6Flip constraints off and the same model free-styles: it may wrap JSON in a markdown7fence, or — more subtly — return valid JSON that violates the schema. That contrast8*is* the demo.9 10Runs on ZeroGPU (H200). No external API, no secrets.11"""12 13import json14import os15import time16 17import gradio as gr18import jsonschema19import spaces20import torch21import outlines22from outlines.types import JsonSchema23from transformers import AutoModelForCausalLM, AutoTokenizer24 25from examples import CONTACT_TEXT, EXAMPLES26from schemas import CUSTOM_LABEL, PRESETS, preset_schema27 28MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-3B-Instruct")29MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "384"))30 31SYSTEM_PROMPT = (32    "You are a precise information-extraction engine. You output a single JSON object and "33    "nothing else. Put each value in the field whose meaning it matches, and never copy the "34    "same value into two different fields."35)36 37DEFAULT_CUSTOM_SCHEMA = json.dumps(38    {39        "type": "object",40        "properties": {41            "summary": {"type": "string"},42            "topics": {"type": "array", "items": {"type": "string"}},43            "sentiment": {"type": "string", "enum": ["positive", "neutral", "negative"]},44        },45        "required": ["summary", "sentiment"],46    },47    indent=2,48)49 50print(f"[init] loading {MODEL_ID} …")51_t0 = time.perf_counter()52_tok = AutoTokenizer.from_pretrained(MODEL_ID)53_hf = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, device_map="cuda")54model = outlines.from_transformers(_hf, _tok)55print(f"[init] model ready in {time.perf_counter() - _t0:.1f}s")56 57 58def resolve_schema(preset: str, custom_schema: str) -> dict:59    if preset == CUSTOM_LABEL:60        return json.loads(custom_schema)61    return preset_schema(preset)62 63 64def build_prompt(text: str, schema: dict) -> str:65    # Naming the fields keeps the model from mis-mapping; the grammar enforces structure.66    fields = ", ".join(schema.get("properties", {}).keys())67    hint = f" with these fields: {fields}" if fields else ""68    user = (69        f"Extract the information from the text below as a JSON object{hint}.\n\n"70        f'Text:\n"""{text}"""\n\n'71        "Return only the JSON object."72    )73    return _tok.apply_chat_template(74        [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user}],75        tokenize=False,76        add_generation_prompt=True,77    )78 79 80@spaces.GPU(duration=60)81def _generate(prompt: str, preset: str, custom_schema: str, constraints_on: bool) -> str:82    # Simple, picklable args only (ZeroGPU forks a worker); rebuild the output type here.83    if not constraints_on:84        return model(prompt, max_new_tokens=MAX_NEW_TOKENS)85    if preset == CUSTOM_LABEL:86        output_type = JsonSchema(json.loads(custom_schema))87    else:88        output_type = PRESETS[preset]89    return model(prompt, output_type=output_type, max_new_tokens=MAX_NEW_TOKENS)90 91 92def extract(text, preset, custom_schema, constraints_on):93    """Extract structured data from free text as JSON that conforms to a schema.94 95    With constraints on, the output is guaranteed valid against the chosen schema96    (right keys, types and enums) via constrained decoding.97 98    Args:99        text: The free text to extract structured information from.100        preset: Which schema to use — "Contact card", "Product", "Job posting",101            "Event", or "Custom (edit the schema)".102        custom_schema: A JSON Schema string; used only when preset is the Custom option.103        constraints_on: If true, force the output to match the schema (recommended).104 105    Returns:106        The extracted JSON (string) and a short validity/status badge (markdown).107    """108    text = (text or "").strip()109    if not text:110        return "", "Paste some text first."111 112    try:113        schema = resolve_schema(preset, custom_schema)114    except json.JSONDecodeError as e:115        return "", f"❌ Your custom schema is not valid JSON: {e}"116 117    prompt = build_prompt(text, schema)118    t0 = time.perf_counter()119    raw = _generate(prompt, preset, custom_schema, constraints_on)120    dt = time.perf_counter() - t0121 122    # 1) Is it even JSON?123    try:124        parsed = json.loads(raw)125    except json.JSONDecodeError as e:126        why = (127            "the model wrapped its answer in a markdown code fence (```)"128            if raw.lstrip().startswith("```")129            else f"`{e}`"130        )131        note = (132            "But constraints were ON — please report this."133            if constraints_on134            else "Constrained decoding never emits a fence or prose — always parseable."135        )136        return raw, f"❌ **Not valid JSON** · {dt:.1f}s — {why}. {note}"137 138    # 2) Does it actually conform to the schema (types, enums, required)?139    try:140        jsonschema.validate(parsed, schema)141    except jsonschema.ValidationError as e:142        shown = json.dumps(parsed, indent=2, ensure_ascii=False)143        note = (144            "But constraints were ON — please report this."145            if constraints_on146            else "Constrained decoding would have *forced* the right type/enum here."147        )148        return shown, f"⚠️ **Valid JSON, but it breaks the schema** · {dt:.1f}s\n\n`{e.message}` — {note}"149 150    pretty = json.dumps(parsed, indent=2, ensure_ascii=False)151    extra = "" if constraints_on else " — the model complied this time, but nothing *forced* it to."152    return pretty, f"✅ **Valid & schema-conformant** · {dt:.1f}s{extra}"153 154 155def on_preset_change(preset):156    return gr.update(visible=(preset == CUSTOM_LABEL))157 158 159INTRO = """160# 🔒 Structured Output Playground161**Lock any LLM's output to a JSON schema.** A local model (Qwen2.5-3B) extracts structured data from162free text. With **constraints ON**, the decoder can only emit tokens that keep the output163**conformant to your schema** — right keys, right *types*, valid *enums*, every time.164 165Flip **OFF** and the same model free-styles: it may wrap the JSON in a markdown fence, or — more166subtly — return *valid JSON that violates your schema* (a string where you asked for an integer, a167value outside your enum). A good model complies *often*; constrained decoding makes it **always**.168 169Runs on **ZeroGPU** (H200) · no external API, no secrets · built by170[Ferr0](https://huggingface.co/Ferr0) · [pixelium.win](https://pixelium.win) · [GitHub](https://github.com/ferr079)171"""172 173with gr.Blocks(title="Structured Output Playground") as demo:174    gr.Markdown(INTRO)175 176    with gr.Row():177        with gr.Column(scale=1):178            preset = gr.Dropdown(179                choices=list(PRESETS.keys()) + [CUSTOM_LABEL],180                value="Contact card",181                label="Schema preset",182            )183            custom = gr.Code(184                value=DEFAULT_CUSTOM_SCHEMA,185                language="json",186                label="Custom JSON Schema",187                visible=False,188            )189            text = gr.Textbox(190                value=CONTACT_TEXT,191                lines=8,192                label="Source text",193                placeholder="Paste any text to extract from…",194            )195            constraints = gr.Checkbox(value=True, label="Constraints ON (force schema)")196            go = gr.Button("Extract", variant="primary")197 198        with gr.Column(scale=1):199            out = gr.Code(label="Extracted JSON", language="json")200            badge = gr.Markdown()201 202    gr.Examples(examples=EXAMPLES, inputs=[text, preset, constraints])203 204    preset.change(on_preset_change, inputs=preset, outputs=custom, api_name=False)205    go.click(extract, inputs=[text, preset, custom, constraints], outputs=[out, badge])206 207 208if __name__ == "__main__":209    demo.launch()210