mlx-lm · Lesson 2

The Python Surface

The same model from Lesson 1, now callable from your own code. Three functions — load, generate, stream_generate — plus the one bridge that makes an Instruct model behave: the chat template. This is the on-ramp to wiring a local model into your own tools.

🎯 Mission step: run LLMs locally — from code, not just the CLI
The one idea

load() once, generate() many. Loading weights into unified memory is the expensive part (~0.8s here, seconds for big models); generation is cheap. So in real code you load a model once at startup and call generate per request — exactly what a server or a pipeline does. The CLI you ran in Lesson 1 is just this, wrapped.

01 Why this is the lesson that unlocks the mission

The CLI is perfect for trying a model. But your mission is to wire models into your own tools — and tools call functions, not subcommands. Every later skill is this surface in disguise: a server is load-then-generate behind HTTP; chat is stream_generate in a loop. Learn these three functions and you can build the rest yourself.

💡 Set up follow-along Everything below lives in one script: ~/projects/learn/public/courses/mlx-lm/practice/lesson02_python.py. Run it with ~/projects/learn/public/courses/mlx-lm/practice/.venv/bin/python ~/projects/learn/public/courses/mlx-lm/practice/lesson02_python.py, or paste lines into a REPL started with ~/projects/learn/public/courses/mlx-lm/practice/.venv/bin/python.

02 The four-line core

This is the whole happy path. Load, turn messages into a prompt, generate:

from mlx_lm import load, generate

model, tokenizer = load("mlx-community/Llama-3.2-1B-Instruct-4bit")   # weights → unified memory

messages = [{"role": "user", "content": "Name three primary colors."}]
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)  # ← the bridge (§03)

text = generate(model, tokenizer, prompt=prompt, max_tokens=40, verbose=False)
print(text)

Real output on your M4 Pro:

# load() returned in 0.79s — type Model + TokenizerWrapper
Here are three primary colors:

1. Red
2. Blue
3. Yellow
💡 What each piece is load(repo) → (model, tokenizer): model runs the forward pass; tokenizer converts text ⇆ token ids and holds the chat template. generate(...) returns the answer as a plain str. verbose=True prints the throughput/memory dashboard you saw in Lesson 1.

03 The chat template — the part everyone gets wrong

An Instruct model wasn't trained on raw text like "Name three primary colors." — it was trained on a specific format with role markers and special tokens. apply_chat_template reproduces that exact format. Here's what your one line of messages actually became (decoded from the token ids it returned):

<|begin_of_text|><|start_header_id|>system<|end_header_id|>

Cutting Knowledge Date: December 2023
Today Date: 18 Jun 2026

<|eot_id|><|start_header_id|>user<|end_header_id|>

Name three primary colors.<|eot_id|><|start_header_id|>assistant<|end_header_id|>

Three things to notice — each is a teachable detail:

Special tokens

<|start_header_id|>, <|eot_id|> etc. mark turn boundaries. The model learned these; plain text has none, so it gets confused.

A system turn appears

Llama injects a default system block — including today's date (18 Jun 2026). The template is code, not a static string.

It ends mid-turn

It stops right after assistant<|end_header_id|> — that's what add_generation_prompt=True does: hand the mic to the assistant so it speaks next.

⚠️ Skip the template and Instruct models misbehave If you pass raw text straight to generate, the model sees an unfamiliar format and tends to ramble, echo your prompt, or never stop cleanly. Rule: for any -Instruct / chat model, always go through apply_chat_template. (Base models — no Instruct — are the exception: feed them raw text to continue.)
💡 One gotcha worth knowing In mlx-lm 0.31.3, apply_chat_template returns a list of token ids (not a string), and generate happily accepts either. To see the formatted text, decode it: tokenizer.decode(prompt) — exactly how the block above was produced.

04 Streaming — tokens as they're produced

generate blocks until the whole answer is ready. For a responsive UI (or a server), you want each token as it lands. stream_generate is a generator that yields a response object per step:

from mlx_lm import load, stream_generate

model, tokenizer = load("mlx-community/Llama-3.2-1B-Instruct-4bit")
messages = [{"role": "user", "content": "Count from 1 to 5."}]
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)

for resp in stream_generate(model, tokenizer, prompt, max_tokens=40):
    print(resp.text, end="", flush=True)   # resp.text = just this step's new text

Each yielded resp carries the new text plus live stats. The fields, straight from your run:

# dir(resp) →
text · token · finish_reason · generation_tokens · generation_tps
prompt_tokens · prompt_tps · peak_memory · logprobs · from_draft

# final step on your machine:
gen tokens: 26 · gen tps: 299.7 · peak memory: 0.793 GB
generatestream_generate
Returnsfinal strgenerator of response objects
You get output…all at once, when donetoken by token, live
Best forbatch jobs, scripts, "just give me the answer"chat UIs, servers, streaming a long answer
Statsprinted if verbose=Trueon every resp (read the last one)
messages [{role, content}] apply_chat_template adds role markers, special tokens → ids model forward pass generate → str stream → tokens
The full Python path: your messages → apply_chat_template turns them into the exact token sequence the model expects → the model decodes → generate hands back a string, or stream_generate hands back tokens one at a time.

05 Steering the output — the sampler

By default mlx-lm is nearly deterministic (it picks high-probability tokens). To make output more varied/creative, pass a sampler built with make_sampler from mlx_lm.sample_utils:

from mlx_lm.sample_utils import make_sampler

sampler = make_sampler(temp=0.8, top_p=0.9)
out = generate(model, tokenizer, prompt=prompt, max_tokens=30, sampler=sampler)
# prompt: "Give me a one-line startup idea."  →  on your machine:
Create a mobile app that uses augmented reality (AR) to transform ordinary
objects into a new, interactive art installation for people of all ages.
KnobWhat it doesTry
tempRandomness. 0 = deterministic/repeatable; higher = more surprising.0.0 for facts/code, 0.7–1.0 for creative
top_pNucleus sampling — only consider the most-likely tokens summing to p.0.9 is a sane default
top_kOnly consider the k most-likely tokens.optional; top_p alone is usually enough
💡 Same idea as the CLI flags mlx_lm.generate --temp 0.8 --top-p 0.9 builds this exact sampler under the hood. CLI flags and Python make_sampler args are the same knobs — one more case of "the CLI is the Python surface, wrapped." More processors live in sample_utils.

06 Check yourself

Think first, then expand. (Stuck? Ask me, your teacher — especially about anything in the chat-template block.)

Q: You load a model and call generate in a loop 100 times. Where does the time go, and what should you not do inside the loop?

A: load() is the expensive one-time cost (weights → memory). generate per call is cheap by comparison. Never call load() inside the loop — load once outside it, reuse model/tokenizer. This is exactly how a server is structured.

Q: An Instruct model is rambling and echoing your prompt back. What's the most likely cause?

A: You passed raw text instead of running it through apply_chat_template. Without the role markers and special tokens it was trained on, the model doesn't recognize "this is a user turn, now answer" — so it just continues text. Wrap your messages with the template.

Q: When would you reach for stream_generate over generate?

A: When you want output as it's produced — a chat UI that types out the answer, a server streaming a response, or any long generation where waiting for the whole thing feels slow. For a batch script that just needs the final string, plain generate is simpler.

Q: You need reproducible output for a test. Which sampler setting?

A: temp=0.0 (greedy decoding) — it always picks the most-likely token, so the same prompt gives the same answer. Raise temp only when you want variety.

07 Going deeper

Keep handy: Glossary & quick reference (now includes chat template, sampler, prefill/decode) · ← Lesson 1. Next up → Lesson 3: mlx_lm.chat — a multi-turn REPL, where context (and the KV cache) starts to matter.