Render MLX chat prompts from standalone Jinja

Author: Boris Peyriguere

complexity/inference/chat_template.py

@@ -6,6 +6,8 @@
 from pathlib import Path
 from typing import Any, Iterable
 
+from jinja2 import Environment, StrictUndefined
+
 CHAT_TEMPLATE_ID = "complexity-chat-v2"
 LEGACY_CHAT_TEMPLATE_ID = "complexity-chat-v1"
 SUPPORTED_CHAT_TEMPLATE_IDS = {LEGACY_CHAT_TEMPLATE_ID, CHAT_TEMPLATE_ID}
@@ -89,6 +91,61 @@ def load_chat_template(path: str | Path | None = None) -> dict[str, Any]:
     return validate_chat_template(json.loads(candidate.read_text(encoding="utf-8")))
 
 
+def load_chat_template_jinja(path: str | Path) -> str:
+    """Load the standalone Jinja contract shipped beside a model bundle."""
+
+    candidate = Path(path)
+    if candidate.is_dir():
+        candidate = candidate / "chat_template.jinja"
+    if not candidate.is_file():
+        raise FileNotFoundError(
+            f"Model bundle has no standalone Jinja chat template: {candidate}"
+        )
+    source = candidate.read_text(encoding="utf-8")
+    if not source.strip():
+        raise ValueError(f"Jinja chat template is empty: {candidate}")
+    return source
+
+
+def render_jinja_messages(
+    messages: Iterable[dict[str, Any]],
+    jinja_source: str,
+    *,
+    eos_token: str,
+    add_generation_prompt: bool,
+) -> str:
+    """Render the same standalone Jinja contract used by HF and vLLM."""
+
+    if not isinstance(eos_token, str) or not eos_token:
+        raise ValueError("A non-empty tokenizer EOS token is required for Jinja rendering")
+    environment = Environment(
+        autoescape=False,
+        undefined=StrictUndefined,
+        keep_trailing_newline=True,
+    )
+    return environment.from_string(jinja_source).render(
+        messages=list(messages),
+        eos_token=eos_token,
+        add_generation_prompt=add_generation_prompt,
+    )
+
+
+def render_jinja_inference_prompt(
+    user_content: str,
+    jinja_source: str,
+    *,
+    eos_token: str,
+) -> str:
+    """Render one user turn and the assistant generation prefix via Jinja."""
+
+    return render_jinja_messages(
+        [{"role": "user", "content": user_content}],
+        jinja_source,
+        eos_token=eos_token,
+        add_generation_prompt=True,
+    )
+
+
 def render_system_prefix(template: dict[str, Any]) -> str:
     system_prompt = str(template["system_prompt"]).strip()
     if not system_prompt:

pyproject.toml

@@ -52,6 +52,7 @@ dependencies = [
     "tqdm>=4.0.0",
     "wandb>=0.15.0",
     "safetensors>=0.4.0",   # save_pretrained + checkpointing
+    "jinja2>=3.0.0",        # standalone HF/vLLM/MLX chat-template rendering
     "pyyaml>=6.0",          # imported at module load in config/model_config.py
     "typer>=0.9.0",         # CLI (complexity.cli)
 ]

scripts/convert_pt_to_mlx.py

@@ -21,6 +21,7 @@
 
 from complexity.inference.chat_template import (
     default_chat_template,
+    huggingface_chat_template,
     validate_chat_template,
 )
 
@@ -186,6 +187,10 @@ def write_chat_template(checkpoint: dict, output: Path) -> dict:
         json.dumps(template, indent=2, sort_keys=True) + "\n",
         encoding="utf-8",
     )
+    (output / "chat_template.jinja").write_text(
+        huggingface_chat_template(template) + "\n",
+        encoding="utf-8",
+    )
     return template
 
 
@@ -213,6 +218,14 @@ def copy_bundle_chat_template(source: Path, output: Path) -> dict:
         json.dumps(template, indent=2, sort_keys=True) + "\n",
         encoding="utf-8",
     )
+    source_jinja = source / "chat_template.jinja"
+    if source_jinja.is_file():
+        shutil.copy2(source_jinja, output / "chat_template.jinja")
+    else:
+        (output / "chat_template.jinja").write_text(
+            huggingface_chat_template(template) + "\n",
+            encoding="utf-8",
+        )
     return template
 
 

scripts/eval_mlx_chat_panel.py

@@ -56,7 +56,11 @@ def generate_one(
 ) -> str:
     rendered = prompt
     if not raw_prompt:
-        rendered, _ = render_bundle_prompt(Path(generation["model_dir"]), prompt)
+        rendered, _ = render_bundle_prompt(
+            Path(generation["model_dir"]),
+            prompt,
+            eos_token=tokenizer.eos_token,
+        )
     mx.random.seed(0)
     sampler = make_sampler(
         temp=float(generation["temperature"]),

scripts/export_tr_hash_vllm.py

@@ -147,7 +147,7 @@ def build_config(raw: dict, chat_template: dict | None = None) -> dict:
         config["use_shared_routed_gates"] = False
     if chat_template is not None:
         config["chat_template_id"] = chat_template["id"]
-        config["chat_template_file"] = "chat_template.json"
+        config["chat_template_file"] = "chat_template.jinja"
     return config
 
 

scripts/mlx_generate.py

@@ -18,16 +18,28 @@
 
 from complexity.inference.chat_template import (
     load_chat_template,
-    render_inference_prompt,
+    load_chat_template_jinja,
+    render_jinja_inference_prompt,
 )
 
 
-def render_bundle_prompt(model_dir: Path, user_message: str) -> tuple[str, str]:
+def render_bundle_prompt(
+    model_dir: Path,
+    user_message: str,
+    *,
+    eos_token: str,
+) -> tuple[str, str]:
     template_path = model_dir / "chat_template.json"
     if not template_path.is_file():
         raise FileNotFoundError(f"MLX bundle has no chat template: {template_path}")
     template = load_chat_template(template_path)
-    return render_inference_prompt(user_message, template), str(template["id"])
+    jinja_source = load_chat_template_jinja(model_dir)
+    rendered = render_jinja_inference_prompt(
+        user_message,
+        jinja_source,
+        eos_token=eos_token,
+    )
+    return rendered, str(template["id"])
 
 
 def main():
@@ -75,7 +87,11 @@ def main():
         tokenizer = load_tokenizer(args.model_dir)
     print(f"Ready. eos_token_id={tokenizer.eos_token_id}\n")
     if args.raw_prompt is None:
-        rendered_prompt, template_id = render_bundle_prompt(args.model_dir, args.prompt)
+        rendered_prompt, template_id = render_bundle_prompt(
+            args.model_dir,
+            args.prompt,
+            eos_token=tokenizer.eos_token,
+        )
         print(f"Chat template: {template_id}")
     else:
         rendered_prompt = args.raw_prompt

tests/test_chat_template.py

@@ -3,14 +3,19 @@
 import json
 from pathlib import Path
 
+import pytest
+
 from complexity.inference.chat_template import (
     CHAT_TEMPLATE_ID,
     LEGACY_CHAT_TEMPLATE_ID,
     THINK_FINAL_ENVELOPE,
     default_chat_template,
     huggingface_chat_template,
+    load_chat_template_jinja,
     render_assistant_envelope,
     render_inference_prompt,
+    render_jinja_inference_prompt,
+    render_jinja_messages,
     render_messages_before_assistant,
     validate_chat_template,
 )
@@ -144,6 +149,43 @@ def test_huggingface_template_matches_explicit_system_message() -> None:
     assert rendered == render_messages_before_assistant(messages, contract)
 
 
+def test_standalone_jinja_renderer_matches_hf_and_vllm_contract(tmp_path) -> None:
+    contract = default_chat_template()
+    source = huggingface_chat_template(contract)
+    (tmp_path / "chat_template.jinja").write_text(source, encoding="utf-8")
+
+    rendered = render_jinja_messages(
+        [
+            {"role": "user", "content": "First"},
+            {"role": "assistant", "content": "Answer"},
+            {"role": "user", "content": "Follow-up"},
+        ],
+        load_chat_template_jinja(tmp_path),
+        eos_token="</s>",
+        add_generation_prompt=True,
+    )
+
+    assert rendered == (
+        "User:\nFirst\n\nAssistant:\nAnswer</s>"
+        "User:\nFollow-up\n\nAssistant:\n"
+    )
+
+
+def test_standalone_jinja_renderer_builds_mlx_generation_prompt() -> None:
+    rendered = render_jinja_inference_prompt(
+        "Hello",
+        huggingface_chat_template(default_chat_template()),
+        eos_token="</s>",
+    )
+
+    assert rendered == "User:\nHello\n\nAssistant:\n"
+
+
+def test_standalone_jinja_loader_rejects_incomplete_mlx_bundle(tmp_path) -> None:
+    with pytest.raises(FileNotFoundError, match="standalone Jinja"):
+        load_chat_template_jinja(tmp_path)
+
+
 def test_vllm_config_declares_exported_template() -> None:
     template = default_chat_template()
     config = build_config(
@@ -158,7 +200,7 @@ def test_vllm_config_declares_exported_template() -> None:
         template,
     )
     assert config["chat_template_id"] == CHAT_TEMPLATE_ID
-    assert config["chat_template_file"] == "chat_template.json"
+    assert config["chat_template_file"] == "chat_template.jinja"
 
 
 def test_vllm_export_preserves_legacy_modulo_cyclic_routing() -> None:
@@ -304,6 +346,9 @@ def test_mlx_export_preserves_checkpoint_chat_template(tmp_path) -> None:
 
     assert written == template
     assert (tmp_path / "chat_template.json").read_text().endswith("\n")
+    assert (tmp_path / "chat_template.jinja").read_text(encoding="utf-8") == (
+        huggingface_chat_template(template) + "\n"
+    )
     assert (
         __import__("json").loads(
             (tmp_path / "chat_template.json").read_text()