hf_chat_template/template.rs
1//! The public [`ChatTemplate`] type and its builder.
2
3use std::sync::Arc;
4
5use minijinja::{Environment, UndefinedBehavior, Value};
6use serde_json::{Map, Value as Json};
7
8use crate::clock::{Clock, SystemClock};
9use crate::config::{self, ChatTemplateField, TokenizerConfig};
10use crate::engine::{self, EngineConfig};
11use crate::error::Error;
12use crate::model::{Message, RenderInput};
13
14/// A compiled Hugging Face chat template, ready to render.
15///
16/// Construction compiles the Jinja source and installs the `transformers`-compatible
17/// globals/filters. [`render`](ChatTemplate::render) borrows `&self`, so a single instance is
18/// cheap to reuse and shareable across threads.
19pub struct ChatTemplate {
20 env: Environment<'static>,
21 /// Special tokens captured from a [`TokenizerConfig`] (empty for raw-string construction),
22 /// injected into the render context as `{{ bos_token }}` etc. unless the input overrides them.
23 special_tokens: Vec<(String, String)>,
24}
25
26impl std::fmt::Debug for ChatTemplate {
27 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28 // The compiled environment is large and not user-meaningful; keep this opaque.
29 f.debug_struct("ChatTemplate")
30 .field("special_tokens", &self.special_tokens)
31 .finish_non_exhaustive()
32 }
33}
34
35impl ChatTemplate {
36 /// Compile from a raw Jinja chat-template string, using the default
37 /// `transformers`-compatible settings (see [`ChatTemplateBuilder`] to customize).
38 ///
39 /// An inherent `from_str` (not just the [`FromStr`](std::str::FromStr) impl) so callers
40 /// can write `ChatTemplate::from_str(s)` without importing the trait.
41 #[allow(clippy::should_implement_trait)]
42 pub fn from_str(source: &str) -> Result<Self, Error> {
43 ChatTemplate::builder(source).build()
44 }
45
46 /// Compile from a parsed `tokenizer_config.json`, resolving the default template and
47 /// injecting the config's special tokens (`bos_token`, …) into the render context.
48 ///
49 /// For named-template selection (`tool_use`, `rag`) or custom options, use
50 /// [`builder_from_config`](ChatTemplate::builder_from_config).
51 pub fn from_tokenizer_config(config: &TokenizerConfig) -> Result<Self, Error> {
52 ChatTemplate::builder_from_config(config)?.build()
53 }
54
55 /// Start a builder to compile `source` with non-default options (clock, undefined policy…).
56 pub fn builder(source: &str) -> ChatTemplateBuilder {
57 ChatTemplateBuilder::new(BuilderSource::Raw(source.to_owned()), Vec::new())
58 }
59
60 /// Start a builder from a `tokenizer_config.json`, carrying its special tokens. Use
61 /// [`template_name`](ChatTemplateBuilder::template_name) to pick a named sub-template.
62 ///
63 /// Fails with [`Error::Config`] if the config has no `chat_template` field at all.
64 pub fn builder_from_config(config: &TokenizerConfig) -> Result<ChatTemplateBuilder, Error> {
65 let field = config
66 .chat_template
67 .clone()
68 .ok_or_else(|| Error::Config("tokenizer config has no chat_template".into()))?;
69 Ok(ChatTemplateBuilder::new(
70 BuilderSource::Field(field),
71 config.special_tokens(),
72 ))
73 }
74
75 /// Fetch `tokenizer_config.json` for a Hub repo (default branch) and compile its default
76 /// template, injecting the config's special tokens. Requires the `hub` feature.
77 ///
78 /// Authentication uses `hf-hub`'s discovery (the `HF_TOKEN` env var or the cached
79 /// `huggingface-cli login` token); gated repos need a token with access. For a pinned
80 /// commit or branch, use [`from_hub_revision`](ChatTemplate::from_hub_revision).
81 ///
82 /// ```no_run
83 /// use hf_chat_template::{ChatTemplate, Message};
84 /// let tmpl = ChatTemplate::from_hub("Qwen/Qwen2.5-0.5B-Instruct")?;
85 /// let prompt = tmpl.render_messages(&[Message::user("Hi")], true)?;
86 /// # Ok::<(), hf_chat_template::Error>(())
87 /// ```
88 #[cfg(feature = "hub")]
89 pub fn from_hub(repo_id: &str) -> Result<Self, Error> {
90 let config = crate::hub::fetch_config(repo_id, None)?;
91 ChatTemplate::from_tokenizer_config(&config)
92 }
93
94 /// Like [`from_hub`](ChatTemplate::from_hub), but pins a specific `revision` (a branch,
95 /// tag, or commit SHA). Requires the `hub` feature.
96 #[cfg(feature = "hub")]
97 pub fn from_hub_revision(repo_id: &str, revision: &str) -> Result<Self, Error> {
98 let config = crate::hub::fetch_config(repo_id, Some(revision))?;
99 ChatTemplate::from_tokenizer_config(&config)
100 }
101
102 /// Render the typed input model to the final prompt string. Special tokens from the
103 /// source config are injected first; any matching key in [`RenderInput::extra`] wins.
104 pub fn render(&self, input: &RenderInput) -> Result<String, Error> {
105 let ctx = self.build_context(input)?;
106 self.render_value(ctx)
107 }
108
109 /// Convenience: render with only `messages` and the generation-prompt flag.
110 pub fn render_messages(
111 &self,
112 messages: &[Message],
113 add_generation_prompt: bool,
114 ) -> Result<String, Error> {
115 let input = RenderInput {
116 messages: messages.to_vec(),
117 add_generation_prompt,
118 ..Default::default()
119 };
120 self.render(&input)
121 }
122
123 /// Render with an arbitrary minijinja context value — the low-level escape hatch for
124 /// callers who need to pass template variables we don't model with typed structs.
125 ///
126 /// Unlike [`render`](ChatTemplate::render), this does **not** inject the config's special
127 /// tokens; the caller's context is used as-is.
128 pub fn render_value(&self, ctx: Value) -> Result<String, Error> {
129 let tmpl = self
130 .env
131 .get_template("chat")
132 .expect("the 'chat' template is always present after construction");
133 tmpl.render(ctx).map_err(Error::from_render)
134 }
135
136 /// Build the Jinja context: special tokens, then the serialized input on top (input wins).
137 /// Routed through `serde_json` to a single `minijinja::Value`; `preserve_order` on both
138 /// crates keeps map key order intact end-to-end (load-bearing for `| tojson`).
139 fn build_context(&self, input: &RenderInput) -> Result<Value, Error> {
140 let mut ctx: Map<String, Json> = Map::new();
141 for (k, v) in &self.special_tokens {
142 ctx.insert(k.clone(), Json::String(v.clone()));
143 }
144 let serialized = serde_json::to_value(input)
145 .map_err(|e| Error::Config(format!("failed to serialize render input: {e}")))?;
146 if let Json::Object(map) = serialized {
147 for (k, v) in map {
148 ctx.insert(k, v);
149 }
150 }
151 Ok(Value::from_serialize(Json::Object(ctx)))
152 }
153}
154
155/// Where a builder draws its template source from.
156enum BuilderSource {
157 /// A raw Jinja string.
158 Raw(String),
159 /// A `chat_template` field whose concrete source is resolved at `build()`.
160 Field(ChatTemplateField),
161}
162
163/// Builder for [`ChatTemplate`], exposing the knobs that affect rendering semantics.
164///
165/// Defaults mirror the `transformers` reference environment:
166/// `trim_blocks = true`, `lstrip_blocks = true`, `keep_trailing_newline = true`,
167/// lenient undefined behavior, pycompat enabled, [`SystemClock`].
168pub struct ChatTemplateBuilder {
169 source: BuilderSource,
170 special_tokens: Vec<(String, String)>,
171 template_name: Option<String>,
172 cfg: EngineConfig,
173}
174
175impl ChatTemplateBuilder {
176 fn new(source: BuilderSource, special_tokens: Vec<(String, String)>) -> Self {
177 ChatTemplateBuilder {
178 source,
179 special_tokens,
180 template_name: None,
181 cfg: EngineConfig {
182 // VERIFY against transformers source; pinned here as the documented baseline.
183 trim_blocks: true,
184 lstrip_blocks: true,
185 keep_trailing_newline: true,
186 undefined: UndefinedBehavior::Lenient,
187 pycompat: true,
188 clock: Arc::new(SystemClock),
189 },
190 }
191 }
192
193 /// Select a named sub-template (e.g. `"tool_use"`, `"rag"`) when the config carries a
194 /// list of them. Ignored for a single-string source.
195 pub fn template_name(mut self, name: impl Into<String>) -> Self {
196 self.template_name = Some(name.into());
197 self
198 }
199
200 /// Inject a deterministic [`Clock`] for `strftime_now` (essential for golden tests).
201 pub fn clock(mut self, clock: impl Clock + 'static) -> Self {
202 self.cfg.clock = Arc::new(clock);
203 self
204 }
205
206 /// Enable/disable the pycompat Python-method shim. Default: enabled.
207 pub fn pycompat(mut self, enabled: bool) -> Self {
208 self.cfg.pycompat = enabled;
209 self
210 }
211
212 /// Override how undefined variables are treated. Default: [`UndefinedBehavior::Lenient`].
213 pub fn undefined_behavior(mut self, ub: UndefinedBehavior) -> Self {
214 self.cfg.undefined = ub;
215 self
216 }
217
218 /// Compile the template. Returns [`Error::Compile`] on a Jinja syntax error, or
219 /// [`Error::Config`] if a requested/needed named template can't be resolved.
220 pub fn build(self) -> Result<ChatTemplate, Error> {
221 let source: String = match &self.source {
222 BuilderSource::Raw(s) => s.clone(),
223 BuilderSource::Field(field) => {
224 config::resolve_template(field, self.template_name.as_deref())?.to_owned()
225 }
226 };
227 let env = engine::build(source, &self.cfg).map_err(Error::Compile)?;
228 Ok(ChatTemplate {
229 env,
230 special_tokens: self.special_tokens,
231 })
232 }
233}
234
235impl std::str::FromStr for ChatTemplate {
236 type Err = Error;
237 fn from_str(s: &str) -> Result<Self, Error> {
238 ChatTemplate::from_str(s)
239 }
240}