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
114pub fn tojson(value: Value, kwargs: Kwargs) -> Result<Value, Error> {
115    if let Ok(indent) = kwargs.get("indent") {
116        let mut buf = Vec::new();
117        let repeat = b" ".repeat(indent);
118        let formatter = serde_json::ser::PrettyFormatter::with_indent(&repeat);
119        let mut serializer = serde_json::Serializer::with_formatter(&mut buf, formatter);
120        value.serialize(&mut serializer).unwrap();
121        String::from_utf8(buf).map_err(|err| {
122            Error::new(ErrorKind::BadSerialization, "cannot serialize to JSON").with_source(err)
123        })
124    } else {
125        serde_json::to_string(&value).map_err(|err| {
126            Error::new(ErrorKind::BadSerialization, "cannot serialize to JSON").with_source(err)
127        })
128    }
129    .map_err(|err| {
130        Error::new(ErrorKind::InvalidOperation, "cannot serialize to JSON").with_source(err)
131    })
132    .map(|s| {
133        // When this filter is used the return value is safe for both HTML and JSON
134        let mut rv = String::with_capacity(s.len());
135        for c in s.chars() {
136            match c {
137                '<' => rv.push_str("\\u003c"),
138                '>' => rv.push_str("\\u003e"),
139                '&' => rv.push_str("\\u0026"),
140                '\'' => rv.push_str("\\u0027"),
141                _ => rv.push(c),
142            }
143        }
144        Value::from_safe_string(rv)
145    })
146}
147
148pub fn strftime_now(format_str: &str) -> Result<Value, Error> {
149    let local: DateTime<Local> = Local::now();
150    Ok(Value::from_safe_string(
151        local.format(format_str).to_string(),
152    ))
153}
dynamo_renderer/template/tokcfg.rs

dynamo_renderer/template/
tokcfg.rs
