Skip to main content

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
179pub fn strftime_now(format_str: &str) -> Result<Value, Error> {
180    let local: DateTime<Local> = Local::now();
181    Ok(Value::from_safe_string(
182        local.format(format_str).to_string(),
183    ))
184}