dynamo_renderer/template/
tokcfg.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//based on: https://github.com/EricLBuehler/mistral.rs/blob/d970bb5feb863acf8e8ec90de97e18221fb959f1/mistralrs-core/src/pipeline/chat_template.rs
5
6use std::collections::HashMap;
7
8use chrono::{DateTime, Local};
9use either::Either;
10use minijinja::{Error, ErrorKind, Value, value::Kwargs};
11use serde::{Deserialize, Serialize};
12
13#[allow(dead_code)]
14#[derive(Debug, Deserialize)]
15pub struct AddedTokensDecoder {
16    __type: Option<String>,
17    pub content: String,
18    lstrip: bool,
19    normalized: bool,
20    rstrip: bool,
21    single_word: bool,
22    special: Option<bool>,
23}
24
25pub fn raise_exception(msg: String) -> Result<String, minijinja::Error> {
26    Err(minijinja::Error::new(ErrorKind::InvalidOperation, msg))
27}
28
29#[derive(Debug, Deserialize)]
30pub struct BeginEndUnkTok(
31    #[serde(with = "either::serde_untagged")] pub Either<String, AddedTokensDecoder>,
32);
33
34/// Support older tool use patterns where the tool use template was separate from the default/chat template.
35/// Modern patterns use a single template with a `tool_use` key, e.g.
36///
37/// ```jinja
38/// {%- if tools is not none and tool_choice is not none %}
39/// ```
40#[derive(Debug, Deserialize)]
41pub struct ChatTemplateValue(
42    #[serde(with = "either::serde_untagged")] pub Either<String, Vec<HashMap<String, String>>>,
43);
44
45/// If present, pad_token is usually a single value. Deepseek R1 and it's distill's use a map.
46#[allow(dead_code)]
47#[derive(Debug, Deserialize)]
48pub struct PadTokenValue(
49    #[serde(with = "either::serde_untagged")] pub Either<String, AddedTokensDecoder>,
50);
51
52#[allow(dead_code)]
53#[derive(Debug, Deserialize, Default)]
54/// Template for chat models including bos/eos/unk as well as the chat template.
55pub struct ChatTemplate {
56    pub bos_token: Option<BeginEndUnkTok>,
57    pub eos_token: Option<BeginEndUnkTok>,
58    pub unk_token: Option<BeginEndUnkTok>,
59
60    /// Jinja format [chat templating] for chat completion.
61    ///
62    /// [chat templating]: https://huggingface.co/docs/transformers/chat_templating
63    pub chat_template: Option<ChatTemplateValue>,
64
65    // future
66    add_bos_token: Option<bool>,
67    add_eos_token: Option<bool>,
68    added_tokens_decoder: Option<HashMap<String, AddedTokensDecoder>>,
69    additional_special_tokens: Option<Vec<String>>,
70    clean_up_tokenization_spaces: Option<bool>,
71    device_map: Option<String>,
72    legacy: Option<bool>,
73    model_max_length: Option<f64>,
74    pad_token: Option<PadTokenValue>,
75    sp_model_kwargs: Option<HashMap<String, String>>,
76    spaces_between_special_tokens: Option<bool>,
77    tokenizer_class: Option<String>,
78    truncation_size: Option<String>,
79    use_default_system_prompt: Option<bool>,
80}
81
82impl ChatTemplate {
83    pub fn eos_tok(&self) -> Option<String> {
84        match self.eos_token.as_ref()?.0 {
85            Either::Left(ref lit) => Some(lit.clone()),
86            Either::Right(ref added) => Some(added.content.clone()),
87        }
88    }
89
90    pub fn bos_tok(&self) -> Option<String> {
91        match self.bos_token.as_ref()?.0 {
92            Either::Left(ref lit) => Some(lit.clone()),
93            Either::Right(ref added) => Some(added.content.clone()),
94        }
95    }
96
97    pub fn unk_tok(&self) -> Option<String> {
98        match self.unk_token.as_ref()?.0 {
99            Either::Left(ref lit) => Some(lit.clone()),
100            Either::Right(ref added) => Some(added.content.clone()),
101        }
102    }
103}
104
105#[allow(dead_code)]
106#[derive(Debug, Deserialize)]
107pub struct GenerationConfig {
108    #[serde(with = "either::serde_untagged")]
109    bos_token_id: Either<u32, Vec<u32>>,
110    #[serde(with = "either::serde_untagged")]
111    eos_token_id: Either<u32, Vec<u32>>,
112}
113
114/// Formatter matching Python `json.dumps` default separators (`", "` and
115/// `": "`). serde_json's `CompactFormatter` writes `","`/`":"` instead, and
116/// chat templates embed these strings directly into the prompt, so the
117/// separator choice is model-visible.
118struct PyJsonFormatter;
119
120impl serde_json::ser::Formatter for PyJsonFormatter {
121    fn begin_array_value<W>(&mut self, writer: &mut W, first: bool) -> std::io::Result<()>
122    where
123        W: ?Sized + std::io::Write,
124    {
125        if !first {
126            writer.write_all(b", ")?;
127        }
128        Ok(())
129    }
130
131    fn begin_object_key<W>(&mut self, writer: &mut W, first: bool) -> std::io::Result<()>
132    where
133        W: ?Sized + std::io::Write,
134    {
135        if !first {
136            writer.write_all(b", ")?;
137        }
138        Ok(())
139    }
140
141    fn begin_object_value<W>(&mut self, writer: &mut W) -> std::io::Result<()>
142    where
143        W: ?Sized + std::io::Write,
144    {
145        writer.write_all(b": ")
146    }
147}
148
149/// Mirrors HF transformers' `tojson` filter, not stock Jinja2's. Transformers
150/// overrides Jinja's HTML-safe `tojson` with plain
151/// `json.dumps(x, ensure_ascii=False)` in its chat-template environment, and
152/// vLLM/SGLang render through that — so chat templates (and the models trained
153/// on their output) expect Python separators and **no** HTML escaping
154/// (`'`, `<`, `>`, `&` stay literal). serde_json leaves non-ASCII unescaped by
155/// default, matching `ensure_ascii=False`.
156pub fn tojson(value: Value, kwargs: Kwargs) -> Result<Value, Error> {
157    let mut buf = Vec::new();
158    let result = if let Ok(indent) = kwargs.get("indent") {
159        // Python `json.dumps(indent=n)` separators are `(",", ": ")` with the
160        // item separator followed by newline + indent — PrettyFormatter matches.
161        let repeat = b" ".repeat(indent);
162        let formatter = serde_json::ser::PrettyFormatter::with_indent(&repeat);
163        let mut serializer = serde_json::Serializer::with_formatter(&mut buf, formatter);
164        value.serialize(&mut serializer)
165    } else {
166        let mut serializer = serde_json::Serializer::with_formatter(&mut buf, PyJsonFormatter);
167        value.serialize(&mut serializer)
168    };
169    result.map_err(|err| {
170        Error::new(ErrorKind::BadSerialization, "cannot serialize to JSON").with_source(err)
171    })?;
172    String::from_utf8(buf)
173        .map_err(|err| {
174            Error::new(ErrorKind::BadSerialization, "cannot serialize to JSON").with_source(err)
175        })
176        .map(Value::from_safe_string)
177}
178
179/// Parse a JSON string into a structured value.
180///
181/// HuggingFace/transformers chat-template environments expose this filter, and several
182/// published templates depend on it — e.g. Step-3.7-Flash's `tool_use` block applies it to
183/// a tool call's `arguments`, which arrive as a JSON *string*, to iterate the decoded
184/// object. Without it minijinja aborts the render with `unknown filter: fromjson`, which
185/// fails every multi-turn tool-call request.
186///
187/// Values that are not strings pass through untouched, so templates that apply the filter
188/// defensively to an already-decoded value keep rendering.
189pub fn fromjson(value: Value) -> Result<Value, Error> {
190    let Some(text) = value.as_str() else {
191        return Ok(value);
192    };
193    let parsed: serde_json::Value = serde_json::from_str(text).map_err(|err| {
194        Error::new(ErrorKind::InvalidOperation, "cannot parse JSON").with_source(err)
195    })?;
196    Ok(Value::from_serialize(&parsed))
197}
198
199pub fn strftime_now(format_str: &str) -> Result<Value, Error> {
200    let local: DateTime<Local> = Local::now();
201    Ok(Value::from_safe_string(
202        local.format(format_str).to_string(),
203    ))
204}
205
206#[cfg(test)]
207mod fromjson_tests {
208    use super::*;
209
210    #[test]
211    fn parses_json_object_string() {
212        let out = fromjson(Value::from(r#"{"location":"San Francisco","n":3}"#)).unwrap();
213        assert_eq!(
214            out.get_attr("location").unwrap().as_str(),
215            Some("San Francisco")
216        );
217        assert_eq!(out.get_attr("n").unwrap().to_string(), "3");
218    }
219
220    #[test]
221    fn parses_json_array_string() {
222        let out = fromjson(Value::from(r#"[1,2,3]"#)).unwrap();
223        assert_eq!(out.len(), Some(3));
224    }
225
226    #[test]
227    fn passes_through_non_string() {
228        // Already-decoded values must survive a defensive `| fromjson`.
229        let already = Value::from_serialize(serde_json::json!({"a": 1}));
230        let out = fromjson(already).unwrap();
231        assert_eq!(out.get_attr("a").unwrap().to_string(), "1");
232    }
233
234    #[test]
235    fn errors_on_malformed_json() {
236        assert!(fromjson(Value::from("{not json")).is_err());
237    }
238
239    /// Regression: renders the shape of Step-3.7-Flash's `tool_use` block, where a tool
240    /// call's `arguments` arrive as a JSON string. Before the filter existed this failed
241    /// with `unknown filter: fromjson`, 500-ing every multi-turn tool-call request.
242    #[test]
243    fn renders_tool_use_block_with_json_string_arguments() {
244        let mut env = minijinja::Environment::new();
245        env.add_filter("fromjson", fromjson);
246        env.add_template(
247            "tool_use",
248            "{% for tc in tool_calls %}{% set a = tc.function.arguments | fromjson %}\
249CALL {{ tc.function.name }} loc={{ a.location }} unit={{ a.unit }}{% endfor %}",
250        )
251        .unwrap();
252        let rendered = env
253            .get_template("tool_use")
254            .unwrap()
255            .render(minijinja::context! { tool_calls => serde_json::json!([{
256                "function": {
257                    "name": "get_weather",
258                    "arguments": "{\"location\":\"San Francisco\",\"unit\":\"F\"}"
259                }
260            }])})
261            .expect("tool_use template must render once `fromjson` is registered");
262        assert_eq!(rendered, "CALL get_weather loc=San Francisco unit=F");
263    }
264}
dynamo_renderer/template/tokcfg.rs

dynamo_renderer/template/
tokcfg.rs
