valeriow/parallel-constrained-decoding

1"""2Hugging Face Spaces Interactive Demo for Parallel Constrained Decoding.3Optimized for Nvidia ZeroGPU (A10G) and PyTorch.4"""5 6import os7import json8import time9from typing import Dict, Any, Generator10 11import gradio as gr12from core.schema import StructuredSchema13from core.engine import run_parallel_generation, run_naive_generation14 15# ZeroGPU decorator support16try:17    import spaces18    gpu_decorator = spaces.GPU(duration=60)19except Exception:20    def gpu_decorator(fn):21        return fn22 23# Load presets from presets/ directory24PRESETS = {}25presets_dir = os.path.join(os.path.dirname(__file__), "presets")26if os.path.exists(presets_dir):27    for fname in sorted(os.listdir(presets_dir)):28        if fname.endswith(".json"):29            try:30                with open(os.path.join(presets_dir, fname), "r") as f:31                    data = json.load(f)32                    title = data.get("title", fname)33                    PRESETS[title] = {34                        "context": data.get("context", ""),35                        "schema": json.dumps(data.get("schema", {}), indent=2)36                    }37            except Exception as e:38                print(f"Error loading {fname}: {e}")39 40preset_titles = list(PRESETS.keys())41default_title = preset_titles[0] if preset_titles else None42default_context = PRESETS[default_title]["context"] if default_title else ""43default_schema = PRESETS[default_title]["schema"] if default_title else "{}"44 45 46@gpu_decorator47def run_comparison(context_str: str, schema_json_str: str):48    if not context_str or not context_str.strip():49        yield (50            "<div style='color: #dc2626; font-weight: 600; padding: 6px 12px;'>Please provide a context prompt.</div>",51            "{}",52            "0.0 ms",53            "{}",54            "0.0 ms"55        )56        return57 58    try:59        schema_dict = json.loads(schema_json_str)60        schema = StructuredSchema(schema_dict)61    except Exception as e:62        yield (63            f"<div style='color: #dc2626; font-weight: 600; padding: 6px 12px;'>Invalid Schema JSON: {e}</div>",64            "{}",65            "0.0 ms",66            "{}",67            "0.0 ms"68        )69        return70 71    try:72        # 1. Run Parallel Constrained Decoding first73        parallel_res = run_parallel_generation(context_str, schema)74        parallel_ms = parallel_res["elapsed_ms"]75        parallel_json_str = json.dumps(parallel_res["parsed_json"], indent=2)76        parallel_time_badge = f"{parallel_ms:.1f} ms"77 78        summary_intermediate = f"""79        <div style="background: #f0fdf4; border: 1px solid #bbf7d0; border-radius: 9999px; padding: 6px 16px; display: inline-flex; align-items: center; gap: 8px; font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; font-size: 14px;">80            <span style="color: #16a34a; font-weight: 700;">Parallel Done: {parallel_time_badge}</span>81            <span style="color: #94a3b8;">·</span>82            <span style="color: #64748b;">Evaluating normal autoregressive baseline...</span>83        </div>84        """85        86        yield (87            summary_intermediate,88            parallel_json_str,89            parallel_time_badge,90            "// Running sequential autoregressive baseline forward passes...",91            "Evaluating..."92        )93 94        # 2. Run Naive generation baseline95        naive_res = run_naive_generation(context_str, schema)96        naive_ms = naive_res["elapsed_ms"]97        naive_json_str = json.dumps(naive_res["parsed_json"], indent=2) if naive_res.get("parsed_json") else naive_res.get("raw_text", "")98        naive_time_badge = f"{naive_ms:.1f} ms"99 100        speedup = round(naive_ms / max(parallel_ms, 1.0), 1)101 102        final_summary_html = f"""103        <div style="background: #f0fdf4; border: 1px solid #bbf7d0; border-radius: 9999px; padding: 8px 20px; display: inline-flex; align-items: center; gap: 10px; font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; font-size: 15px; box-shadow: 0 1px 3px rgba(0,0,0,0.05);">104            <span style="color: #16a34a; font-weight: 800; font-size: 16px; letter-spacing: 0.5px;">{speedup}x FASTER</span>105            <span style="color: #cbd5e1; font-weight: 600;">·</span>106            <span style="color: #334155; font-weight: 600; font-family: monospace;">{parallel_time_badge} vs {naive_time_badge}</span>107        </div>108        """109 110        yield (111            final_summary_html,112            parallel_json_str,113            parallel_time_badge,114            naive_json_str,115            naive_time_badge116        )117    except Exception as err:118        import traceback119        err_msg = f"{err}\n{traceback.format_exc()}"120        yield (121            f"<div style='color: #dc2626; background: #fef2f2; border: 1px solid #fecaca; border-radius: 8px; padding: 10px 14px; font-family: monospace; font-size: 13px;'>Error: {err}</div>",122            "{}",123            "0.0 ms",124            f"Error details:\n{err_msg}",125            "0.0 ms"126        )127 128 129with gr.Blocks(title="Parallel Constrained Decision Engine") as demo:130    gr.Markdown("# Parallel Constrained vs Normal Inference (Qwen2.5 1.5B)")131    gr.Markdown("Parallel Constrained Decoding evaluates all schema fields simultaneously against broadcast prefix KV-cache states, delivering substantial latency reductions with 100% schema adherence.")132 133    with gr.Row():134        preset_dropdown = gr.Dropdown(135            choices=preset_titles,136            value=default_title,137            label="Select Preset Scenario",138            scale=4139        )140        btn_run = gr.Button("⚡ Run Comparison", variant="primary", scale=1)141 142    summary_banner = gr.HTML(value="")143 144    with gr.Row():145        with gr.Column(scale=1):146            gr.Markdown("### Parallel Constrained (Qwen2.5 1.5B)")147            timer_parallel = gr.Textbox(label="Elapsed Time", value="0.0 ms", interactive=False, max_lines=1)148            output_parallel = gr.Code(label="Parallel JSON (Values + Calibrated Probabilities)", language="json", interactive=False, lines=18)149        150        with gr.Column(scale=1):151            gr.Markdown("### Normal Inference (Qwen2.5 1.5B)")152            timer_naive = gr.Textbox(label="Elapsed Time", value="0.0 ms", interactive=False, max_lines=1)153            output_naive = gr.Code(label="Autoregressive JSON Output", language="json", interactive=False, lines=18)154 155    with gr.Accordion("Inspect Context Document & Schema Definition", open=False):156        context_input = gr.Textbox(157            label="Context Document",158            value=default_context,159            lines=6160        )161        schema_input = gr.Code(162            label="Schema Definition (JSON)",163            value=default_schema,164            language="json",165            lines=10166        )167 168    def on_preset_change(title):169        if title in PRESETS:170            return PRESETS[title]["context"], PRESETS[title]["schema"]171        return "", "{}"172 173    preset_dropdown.change(174        fn=on_preset_change,175        inputs=[preset_dropdown],176        outputs=[context_input, schema_input]177    )178 179    btn_run.click(180        fn=run_comparison,181        inputs=[context_input, schema_input],182        outputs=[summary_banner, output_parallel, timer_parallel, output_naive, timer_naive]183    )184 185if __name__ == "__main__":186    demo.queue().launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))187