from mlx_lm import load, generate

model, tokenizer = load("omercelik/Trace-Inverter-4B-MLX-8bit")

messages = [
    {
        "role": "system",
        "content": (
            "You are a trace inversion model. Given a problem, a final answer, "
            "and several compressed reasoning bubbles, reconstruct a detailed "
            "reasoning trace that could plausibly lead to the final answer."
        ),
    },
    {
        "role": "user",
        "content": """Problem:
If a pizza needs 10 cups of water, 16 cups of flour, and salt equal to half the flour amount, what is the combined total?

Model final answer:
34 cups.

Reasoning Bubbles:
I need to calculate the salt first because it is defined as half of the flour amount. Then I should add water, flour, and salt together to get the combined total.

Reconstruct the full reasoning trace.""",
    },
]

prompt = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_dict=False,
)

response = generate(
    model,
    tokenizer,
    prompt=prompt,
    max_tokens=512,
    verbose=True,
)

Notes

The source checkpoint stores PEFT-style LoRA-wrapped tensors inside the safetensors files. For MLX compatibility, the LoRA tensors were merged into plain model weights before conversion. The inferred LoRA scale used for the merge was 1.0.

The source model card notes that outputs may occasionally include stray tool tags such as <tool_call>. Post-processing is recommended when generating datasets.

Generated traces are synthetic reasoning traces. They should not be treated as recovered hidden chain-of-thought from any closed model.