mlx-lm · Lesson 8

Structured Output — JSON You Can Trust

Lesson 7 ended on a problem: the 1B tagger drifted out of format. Here's the fix that turns a small local model into a reliable component — constrain the decoder to a schema so it cannot produce invalid output. Not a better prompt; a mathematical guarantee.

🎯 Mission step: make local-model output reliable enough to parse in code
The one idea

Don't ask for JSON and hope — mask every token that would break it. Compile your schema into a finite-state machine, and at each decoding step set the logits of all illegal tokens to −∞. The model literally cannot emit a character that violates the schema. Output that always parses, every time, from a 1B model.

01 Why "prompt and hope" isn't enough

In Lesson 7 your tagger sometimes returned Category: Security | Summary: … and sometimes wandered into prose. For a human that's fine; for code that does json.loads(output), one stray token is a crash. The smaller the model, the more it drifts — and small local models are exactly what you want for cheap, high-volume work. So you need structural reliability, not a sterner prompt.

💡 Two ways to get structure Prompt-and-parse: ask nicely, then validate/retry — probabilistic, fails on small models. Constrained decoding: restrict what tokens are even possible at each step so the output is valid by construction — deterministic guarantee. This lesson does the second.

02 How constrained decoding works

A model picks each token from a probability distribution over the whole vocabulary. Constrained decoding inserts a step: given "what's been written so far," a finite-state machine (built from your schema) knows which tokens keep the output valid. Everything else gets masked to −∞ before sampling.

so far: {"category": vocabulary logits "Security "AI & ML Sure! 42 → FSM mask legal tokens only illegal → logit −∞ masked (would break schema) sample from these
Only tokens that keep the output on a valid path through the schema's FSM survive; the model samples from those. A category field constrained to a fixed list can only emit one of those exact strings — which is also why the 1B model below finally classifies MLX correctly.

03 The code — your schema is the contract

The Outlines library has an mlx-lm backend. Install it into your practice venv, define a Pydantic model, and pass it as output_type:

# uv pip install "outlines[mlxlm]"   (per your uv preference)
import outlines, mlx_lm, json
from pydantic import BaseModel, Field
from typing import Literal, List

class Tag(BaseModel):
    category: Literal["AI & ML", "Programming Languages", "Security", "Other"]
    summary:  str       = Field(max_length=80)     # bound it (see §05)
    keywords: List[str] = Field(max_length=5)

model = outlines.from_mlxlm(*mlx_lm.load("mlx-community/Llama-3.2-1B-Instruct-4bit"))

out = model("Classify and summarize: A SQL injection smuggles commands into a query.",
            output_type=Tag, max_tokens=256)
tag = Tag.model_validate_json(out)        # never raises — out is valid by construction

Real output on your M4 Pro — three inputs, all parsed & validated, zero format errors:

# category                  keywords
Security                ['SQL Injection', 'Attack', 'Smuggling', 'Commands', 'Query']
Programming Languages   ['React', 'useEffect', 'side effects', 'rendering']
AI & ML                 ['Apple Silicon', 'MLX', 'LLMs', 'Large Language Models']
🎯 Your tangible win Run ~/projects/learn/public/courses/mlx-lm/practice/.venv/bin/python ~/projects/learn/public/courses/mlx-lm/practice/structured_tagger.py. The same 1B model that drifted in Lesson 7 now returns JSON that model_validate_json accepts every time — and it even gets MLX → AI & ML right, because category can only be one of the allowed literals.

04 Guaranteed structure ≠ guaranteed completion

⚠️ A gotcha worth internalizing (hit live while building this) The FSM guarantees the output is a valid prefix of your schema at every step — but if you cut generation off with max_tokens before it closes the JSON, you get a truncated (still-invalid-overall) string. First attempts here failed with Unterminated string because the model wrote a long summary and ran out of budget. Two fixes, both shown above: bound free-text fields in the schema (Field(max_length=80)) so they must close, and give a sane max_tokens. Constrain the shape and the size.
Prompt & parseConstrained decoding
Validityprobabilistic — may failguaranteed by construction
Small modelsdrifts, needs retriesworks — the model can't drift
Enumsmay invent valuesLiteral[...] → only allowed values
Costretries on failureone-time FSM compile + tiny per-token mask
Caveat—not for batched generation; still budget tokens

05 Where this changes what you can build

Reliable structure is what makes a local model a component instead of a toy. With guaranteed JSON you can safely build: classifiers/routers, data extraction from messy text, tool-calling (the model emits a valid function-call object), form-filling, and grading/judging pipelines — all on-device, all parseable without defensive try/except gymnastics.

💡 This retro-fixes Lesson 7 Your Lesson 7 capstone tagger was the right pattern; this is the reliability layer it was missing. Swap the free-form generate for a constrained model(..., output_type=Schema) and the tagger becomes production-grade — exactly what you'd want before wiring it into the learn pipeline.

06 Check yourself

Think first, then expand. (Want a schema for a specific tool — a function-call object, an extraction shape? Ask me and we'll design it.)

Q: Why can't the constrained model return "category": "Frontend" if your Literal only lists four values?

A: At the step where it writes the category value, the FSM only permits tokens that continue one of the four allowed strings. "Frontend" isn't a valid path, so those tokens are masked to −∞ and can't be sampled. The enum is enforced, not requested.

Q: You constrained the schema but still got Unterminated string. What happened and how do you fix it?

A: Generation hit max_tokens before the JSON closed — valid-so-far but incomplete. Bound free-text fields with Field(max_length=N) and/or raise max_tokens so it has room to finish.

Q: When is constrained decoding overkill?

A: When you want free-form prose (a chat reply, a summary a human reads), or when a big model already nails the format and you're not parsing it in code. Use it where a program consumes the output.

Q: Does this make a small model as accurate as a big one?

A: No — it guarantees format, not correctness. The model can still pick the wrong (but valid) category. Constraining the enum helps, but for higher accuracy you still reach for a bigger model or a fine-tune.

07 Going deeper

Keep handy: Glossary & quick reference (see constrained decoding, Outlines) · ← Lesson 7. Next up → Lesson 9: Model Selection & Memory Math — how to know in advance what actually fits and runs well on your M4 Pro (params × bits → GB, plus KV-cache overhead).