Hugging Face's logo

Spaces:
drinkmoonshine
/
parallel-constrained-decoding
Running on Zero

App Files Community
parallel-constrained-decoding
File size: 5,065 Bytes
a0270e2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
---
language:
- en
license: apache-2.0
library_name: mlx
tags:
- structured-generation
- parallel-decoding
- constrained-decoding
- apple-silicon
- mlx
- classification
- json
pipeline_tag: text-generation
base_model: Qwen/Qwen2.5-1.5B-Instruct
---

# Qwen2.5-1.5B-Instruct with Parallel Constrained Decoding

This repository provides an inference implementation for structured JSON generation and high-cardinality classification using `mlx-community/Qwen2.5-1.5B-Instruct-4bit` on Apple Silicon.

Instead of generating structured JSON token-by-token through sequential autoregressive loops, this engine uses **Parallel Constrained Decoding**. It broadcasts the model KV-cache across all schema fields simultaneously, evaluating all decisions in parallel forward passes.

## Key Performance Highlights (Apple Silicon M4 Max)

- **High-Cardinality Decisions (255 choices)**: 89 ms total latency vs. 500 ms autoregressive baseline (5.6x faster).
- **Enterprise Multi-Field Extraction (28 fields)**: 270 ms total latency vs. 1,900 ms autoregressive baseline (7.0x faster).
- **Guaranteed Schema Validity**: 100% valid JSON syntax without grammar parsers, rejection sampling, or repair loops.
- **Calibrated Field Confidence**: Exact softmax probabilities computed directly over candidate token logits for every field.
- **Unified Memory Footprint**: ~1.1 GB total RAM footprint in 4-bit quantization on Apple Silicon.

## How It Works

Traditional structured output engines run standard autoregressive decoding. For an N-field JSON schema, the model performs hundreds of sequential forward passes:

```
[System + Prompt] -> Token 1 -> Token 2 -> ... -> Token K (O(N) sequential forward passes)
```

Parallel Constrained Decoding decomposes the structured generation task into an isolated broadcast pass:

1. **Prefix Prefill**: The context and semantic schema instructions are prefilled once. The resulting Key-Value (KV) cache is held in Apple Silicon Unified Memory.
2. **KV-Cache Broadcasting**: The KV-cache is broadcast across all target fields concurrently.
3. **Sub-Vocabulary Projection**: For each field, only valid candidate choices (e.g. enum options or boolean states) are evaluated. Unrelated vocabulary tokens are masked out.
4. **Calibrated Softmax**: Probabilities are computed directly via softmax over the candidate logit slice:
   $$P(c_i) = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}$$
5. **Collision Disambiguation**: In cases where candidate tokens share prefix strings, the engine follows continuation slices with zero memory reallocation.
6. **Programmatic Assembly**: The verified field choices and confidence scores are formatted directly into structured JSON.

## Quickstart SDK

### Installation

```bash
pip install -r requirements.txt
```

### Python Usage

```python
from core.schema import StructuredSchema
from core.engine import run_parallel_generation

# 1. Define schema
schema_definition = {
    "fraud_risk": {
        "type": "enum",
        "choices": ["LOW", "ELEVATED", "SUSPICIOUS", "CRITICAL"],
        "description": "Risk assessment tier for incoming transaction"
    },
    "block_account": {
        "type": "boolean",
        "description": "Whether immediate account restriction is required"
    },
    "recommended_action": {
        "type": "enum",
        "choices": ["ALLOW", "STEP_UP_2FA", "TEMPORARY_HOLD", "TERMINATE_SESSION"],
        "description": "Immediate mitigation action"
    }
}

schema = StructuredSchema(schema_definition)

# 2. Provide context
context = """
User ID: usr_9921
Location: Lagos, Nigeria (usual: Seattle, USA)
Device: Unknown Linux Chromium browser
Action: Wire transfer $49,500 to offshore escrow
Prior velocity: 0 transfers in 90 days
"""

# 3. Execute parallel generation
result = run_parallel_generation(context, schema)

print(f"Elapsed Time: {result['elapsed_ms']} ms")
print(f"Sequential Passes: {result['sequential_forward_passes']}")
print(f"Parsed JSON: {result['parsed_json']}")
```

### Output Example

```json
{
  "fraud_risk": { "value": "CRITICAL", "prob": 0.9942 },
  "block_account": { "value": "true", "prob": 0.9881 },
  "recommended_action": { "value": "TEMPORARY_HOLD", "prob": 0.9715 }
}
```

## Model Details

- **Base Model**: Qwen/Qwen2.5-1.5B-Instruct
- **Quantization**: 4-bit AWQ (mlx-community format)
- **Context Window**: 32,768 tokens
- **Hardware Target**: Apple Silicon (M1, M2, M3, M4 series with unified memory)
- **Supported Field Types**: Categorical Enums (up to 255 choices per field) and Booleans

## Benchmark Summary

Evaluated on Apple Silicon M4 Max (128GB Unified Memory, MLX 0.22+):

| Scenario | Schema Fields | Autoregressive (ms) | Parallel Constrained (ms) | Speedup | Valid Syntax |
| :--- | :--- | :--- | :--- | :--- | :--- |
| Fintech Fraud Routing | 4 fields | 420 ms | 75 ms | **5.6x** | 100% |
| Code Security Audit | 4 fields | 380 ms | 68 ms | **5.6x** | 100% |
| High-Cardinality Tariff | 1 field (255 choices) | 500 ms | 89 ms | **5.6x** | 100% |
| Support Triage Matrix | 28 fields | 1,900 ms | 270 ms | **7.0x** | 100% |