Support for structured output in batch mode?

[do-me]

I'm running entity matching locally on my Mac and heavily rely on structured output to get the best results. At the same time, I don't want to compromise on performance and just found that the batch_generate function does not yet support structured output. Would be cool if this could be implemented at some point so that mlx-lm can become a real workhorse.

[lazarust]

Looking into this. Should have a PR in the next week or so.

[lazarust]

@do-me Do you have an example script of how you're wanting to use this and what the expected output would be?

[do-me]

Sure. High-level it's fairly straight-forward:

I have a labeled dataset (parquet table) where in each row I'm comparing entity A (e.g. "Starbucks Café) against entity B (e.g. "Starbucks Coffee"), 0 is no match, 1 is match. A human labeled it but I suspected that an LLM can do a similarly good job.
Now I ran some experiments with Qwen/Qwen3-next-80b to confirm my hypothesis and it turned out that I get an alignment of >0.8 Cohens Kappa so excellent human alignment (corresponds to more than 0.9 raw alignment).

I found that you can pass a low one-digit number of rows like e.g. 2 or 3 at the same time without a huge drop in F1 and accuracy compared to just passing just 1 entity at a time. Still there is a drop in F1 and accuracy. As this labeled dataset will be used for training I need to max out F1 and accuracy so I can only pass one row at a time.

Or in simple terms: 1k rows -> 1k prompts.

I use structured output to get the results properly formatted, three keys per row:

• match true/false
• confidence 0-1
• reasoning 2 sentences

I did some benchmarks in the past (willccbb/mlx_parallm#5) where I found significant speedups in batch modes and some weirder things, that e.g. running 31 prompts take as long as running 6 prompts (in batch mode).

Anyway, long story short, here is a very simplified script (not tested):

import json
import time
from mlx_lm import load, generate, batch_generate

# 1. mock dataset (POI Pairs)
# In reality, this comes from a parquet file, but here we hardcode 3 examples.
mock_data = [
    {"id": "a1", "target": "Starbucks, 5th Ave", "base": "Starbucks Coffee, 5th Avenue", "label": 1},
    {"id": "b2", "target": "McDonalds, Main St", "base": "Burger King, Main St", "label": 0},
    {"id": "c3", "target": "Unknown Shop", "base": "Shop Unknown", "label": 2},
]

# 2. structured output to enforce
schema = {
    "type": "object",
    "properties": {
        "pair_index": {"type": "integer"},
        "answer": {"type": "integer", "enum": [0, 1, 2]},
        "confidence": {"type": "number"},
        "reasoning": {"type": "string"}
    },
    "required": ["pair_index", "answer", "confidence", "reasoning"],
    "additionalProperties": False
}

# 3. basic prompts, more complex in reality
system_prompt = "You are an entity resolution expert. Output valid JSON only."
prompts = []

for i, row in enumerate(mock_data):
    user_content = f"""
    Compare these two entities:
    Entity A: {row['target']}
    Entity B: {row['base']}
    
    Return JSON with:
    - pair_index: {i}
    - answer: 0 (No Match), 1 (Match), 2 (Unsure)
    - confidence: 0.0 to 1.0
    - reasoning: short explanation
    """
    
    # Format for chat model
    prompts.append(
        f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
        f"<|im_start|>user\n{user_content}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )

# 4. LOAD MODEL
model_path = "qwen/qwen3-next-80b" 
print(f"Loading {model_path}...")
model, tokenizer = load(model_path)

# ==============================================================================
# SCENARIO A: Currenty slow as sequential 
# ==============================================================================
print("\n--- SCENARIO A: Sequential `generate` (Supports Schema) ---")
start = time.time()
for prompt in prompts:
    # This works perfectly but is slow for thousands of rows
    response = generate(
        model, 
        tokenizer, 
        prompt=prompt, 
        max_tokens=1024, 
        verbose=False,
        json_schema=schema # <--- working!
    )
    print(f"Gen Output: {response[:50]}...") 
print(f"Time taken: {time.time() - start:.2f}s")

# ==============================================================================
# SCENARIO B: batch request with schema
# ==============================================================================
print("\n--- SCENARIO B: `batch_generate` (Feature Request) ---")
print("We want to pass `json_schema` here to enforce structure across the batch.")

start = time.time()
# runs fast but output structure is not guaranteed/enforced
responses = batch_generate(
    model, 
    tokenizer, 
    prompts, 
    max_tokens=200, 
    verbose=False,
    # json_schema=schema  <--- this is missing
)
print(f"Batch Output count: {len(responses)}")
print(f"First output: {responses[0][:50]}...")
print(f"Time taken: {time.time() - start:.2f}s")

[lazarust]

I've got a couple of ideas on how to implement, but I'm unsure which is the best.

Ideas:

1. Handle the JSON enforcement in a logits processor so the schema is enforced at the time of generation
2. Add the schema to the prompt similar to how this works for generating using DeepSeek 3.2

Any input on which direction to head down?

Also tagging @awni, in case they have any opinions.