|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +"""Render OpenAI chat messages into a single prompt string. |
| 8 | +
|
| 9 | +The ExecuTorch runner tokenizes a plain prompt; chat formatting is the server's |
| 10 | +job (control plane). We require the model's own Hugging Face ``chat_template`` |
| 11 | +(via ``--hf-tokenizer``) for correct, tool-aware, reasoning-aware formatting. |
| 12 | +The generic ChatML fallback is opt-in only (``allow_fallback``): it is |
| 13 | +approximate and cannot reproduce model-specific controls (e.g. enable_thinking), |
| 14 | +so it must be a deliberate choice rather than a silent default. |
| 15 | +""" |
| 16 | + |
| 17 | +import json |
| 18 | +import logging |
| 19 | +from typing import Any, Optional |
| 20 | + |
| 21 | +from .protocol import ChatMessage |
| 22 | + |
| 23 | +logger = logging.getLogger(__name__) |
| 24 | + |
| 25 | + |
| 26 | +_DEFAULT_SPECIAL_TOKENS = ["<|im_end|>", "<|endoftext|>", "<|eot_id|>", "<|end|>"] |
| 27 | + |
| 28 | + |
| 29 | +def _decode_tool_call_arguments(messages: list[dict[str, Any]]) -> None: |
| 30 | + """In-place: parse each tool call's ``function.arguments`` from a JSON string |
| 31 | + into an object. |
| 32 | +
|
| 33 | + OpenAI sends assistant tool-call arguments as a JSON-encoded string, but HF |
| 34 | + chat templates expect a mapping (e.g. Qwen renders ``arguments|items`` into |
| 35 | + ``<parameter=…>`` tags). Without this, a multi-turn tool conversation makes |
| 36 | + the template raise "Can only get item pairs from a mapping". Left as-is if |
| 37 | + the value isn't valid JSON, so a template that wants the raw string still works. |
| 38 | + """ |
| 39 | + for m in messages: |
| 40 | + for tc in m.get("tool_calls") or []: |
| 41 | + fn = tc.get("function") |
| 42 | + if not isinstance(fn, dict): |
| 43 | + continue |
| 44 | + args = fn.get("arguments") |
| 45 | + if isinstance(args, str): |
| 46 | + try: |
| 47 | + fn["arguments"] = json.loads(args) |
| 48 | + except (ValueError, TypeError): |
| 49 | + pass |
| 50 | + |
| 51 | + |
| 52 | +class ChatTemplate: |
| 53 | + def __init__( |
| 54 | + self, |
| 55 | + hf_tokenizer_path: Optional[str] = None, |
| 56 | + default_template_kwargs: Optional[dict[str, Any]] = None, |
| 57 | + allow_fallback: bool = False, |
| 58 | + ): |
| 59 | + # Server-level defaults (e.g. {"enable_thinking": False}); per-request |
| 60 | + # chat_template_kwargs override these. |
| 61 | + self._defaults = default_template_kwargs or {} |
| 62 | + self._hf = None |
| 63 | + if hf_tokenizer_path: |
| 64 | + from transformers import AutoTokenizer |
| 65 | + |
| 66 | + self._hf = AutoTokenizer.from_pretrained(hf_tokenizer_path) |
| 67 | + if self._hf.chat_template is None: |
| 68 | + self._hf = None |
| 69 | + if not allow_fallback: |
| 70 | + raise ValueError( |
| 71 | + f"HF tokenizer at {hf_tokenizer_path} has no chat_template; " |
| 72 | + "pass an explicit fallback flag to use approximate ChatML." |
| 73 | + ) |
| 74 | + logger.warning( |
| 75 | + "No chat_template at %s; using approximate ChatML.", |
| 76 | + hf_tokenizer_path, |
| 77 | + ) |
| 78 | + elif not allow_fallback: |
| 79 | + raise ValueError( |
| 80 | + "A chat template is required: pass --hf-tokenizer for the model's own " |
| 81 | + "template, or opt into approximate ChatML with --allow-chatml-fallback." |
| 82 | + ) |
| 83 | + else: |
| 84 | + logger.warning( |
| 85 | + "No --hf-tokenizer; using approximate ChatML (no thinking control)." |
| 86 | + ) |
| 87 | + |
| 88 | + def render( |
| 89 | + self, |
| 90 | + messages: list[ChatMessage], |
| 91 | + tools: Optional[list[dict[str, Any]]] = None, |
| 92 | + template_kwargs: Optional[dict[str, Any]] = None, |
| 93 | + ) -> str: |
| 94 | + kwargs = {**self._defaults, **(template_kwargs or {})} |
| 95 | + if self._hf is not None: |
| 96 | + dumped = [m.model_dump(exclude_none=True) for m in messages] |
| 97 | + _decode_tool_call_arguments(dumped) |
| 98 | + return self._hf.apply_chat_template( |
| 99 | + dumped, |
| 100 | + tools=tools, |
| 101 | + add_generation_prompt=True, |
| 102 | + tokenize=False, |
| 103 | + **kwargs, |
| 104 | + ) |
| 105 | + return self._fallback(messages) |
| 106 | + |
| 107 | + def chat_template_str(self) -> Optional[str]: |
| 108 | + """Raw chat-template string (for tool-format auto-detection), if available.""" |
| 109 | + return ( |
| 110 | + getattr(self._hf, "chat_template", None) if self._hf is not None else None |
| 111 | + ) |
| 112 | + |
| 113 | + def count_tokens(self, prompt: str) -> Optional[int]: |
| 114 | + """Token count for the rendered prompt, or None if no tokenizer is available.""" |
| 115 | + if self._hf is not None: |
| 116 | + # The prompt is already rendered (apply_chat_template includes the |
| 117 | + # control tokens), so encode without re-adding BOS/EOS — matching the |
| 118 | + # session/prefix-cache paths, so the count isn't inflated and |
| 119 | + # near-limit requests aren't falsely rejected under --max-context. |
| 120 | + return len(self._hf.encode(prompt, add_special_tokens=False)) |
| 121 | + return None |
| 122 | + |
| 123 | + def special_tokens(self) -> list[str]: |
| 124 | + """Special-token strings whose appearance ends the visible content. |
| 125 | +
|
| 126 | + From the HF tokenizer when available (model-accurate), else a default set |
| 127 | + covering common chat models. |
| 128 | + """ |
| 129 | + if self._hf is not None: |
| 130 | + toks = list(getattr(self._hf, "all_special_tokens", []) or []) |
| 131 | + return [t for t in toks if isinstance(t, str) and t] |
| 132 | + return list(_DEFAULT_SPECIAL_TOKENS) |
| 133 | + |
| 134 | + @staticmethod |
| 135 | + def _fallback(messages: list[ChatMessage]) -> str: |
| 136 | + # Approximate ChatML. Provide --hf-tokenizer for model-correct formatting |
| 137 | + # (including reasoning controls like enable_thinking, which the fallback |
| 138 | + # cannot reproduce). |
| 139 | + parts = [] |
| 140 | + for m in messages: |
| 141 | + content = m.content if isinstance(m.content, str) else str(m.content or "") |
| 142 | + parts.append(f"<|im_start|>{m.role}\n{content}<|im_end|>") |
| 143 | + parts.append("<|im_start|>assistant\n") |
| 144 | + return "\n".join(parts) |
0 commit comments