Skip to main content

hf_chat_template/
config.rs

1//! `tokenizer_config.json` loading and template-source resolution.
2//!
3//! A [`TokenizerConfig`] is a tolerant view over the parts of a model's `tokenizer_config.json`
4//! we need: the `chat_template` (in any of its historical shapes) plus the special-token fields
5//! the templates reference as context variables (`{{ bos_token }}`). Unknown fields are kept in
6//! `extra`, never rejected.
7
8use serde::{Deserialize, Serialize};
9use serde_json::{Map, Value as Json};
10
11use crate::error::Error;
12
13/// A tolerant view over a parsed `tokenizer_config.json`.
14#[derive(Clone, Debug, Default, Serialize, Deserialize)]
15pub struct TokenizerConfig {
16    /// The chat template, in whichever shape this config uses (see [`ChatTemplateField`]).
17    #[serde(default, skip_serializing_if = "Option::is_none")]
18    pub chat_template: Option<ChatTemplateField>,
19
20    /// Beginning-of-sequence token, if the model defines one.
21    #[serde(default, skip_serializing_if = "Option::is_none")]
22    pub bos_token: Option<TokenField>,
23    /// End-of-sequence token.
24    #[serde(default, skip_serializing_if = "Option::is_none")]
25    pub eos_token: Option<TokenField>,
26    /// Padding token.
27    #[serde(default, skip_serializing_if = "Option::is_none")]
28    pub pad_token: Option<TokenField>,
29    /// Unknown token.
30    #[serde(default, skip_serializing_if = "Option::is_none")]
31    pub unk_token: Option<TokenField>,
32
33    /// Everything else in the file, preserved.
34    #[serde(default, flatten)]
35    pub extra: Map<String, Json>,
36}
37
38/// The `chat_template` field's three historical shapes.
39#[derive(Clone, Debug, Serialize, Deserialize)]
40#[serde(untagged)]
41pub enum ChatTemplateField {
42    /// A single Jinja string: `"chat_template": "{% for … %}"`.
43    Single(String),
44    /// A list of named templates: `[{"name":"default","template":"…"}, …]`.
45    Named(Vec<NamedTemplate>),
46}
47
48/// One entry of the list-of-named-templates form.
49#[derive(Clone, Debug, Serialize, Deserialize)]
50pub struct NamedTemplate {
51    /// The template's name (e.g. `"default"`, `"tool_use"`, `"rag"`).
52    pub name: String,
53    /// The Jinja source.
54    pub template: String,
55}
56
57/// A special-token field: either a bare string, or an `AddedToken` object with a `content` key.
58#[derive(Clone, Debug, Serialize, Deserialize)]
59#[serde(untagged)]
60pub enum TokenField {
61    /// `"bos_token": "<s>"`.
62    Str(String),
63    /// `"bos_token": {"content": "<s>", "lstrip": false, …}`.
64    Obj(Map<String, Json>),
65}
66
67impl TokenField {
68    /// The token's string value (`content` for the object form).
69    pub fn as_str(&self) -> Option<&str> {
70        match self {
71            TokenField::Str(s) => Some(s),
72            TokenField::Obj(m) => m.get("content").and_then(Json::as_str),
73        }
74    }
75}
76
77impl TokenizerConfig {
78    /// The special tokens to inject into the render context, as `(name, value)` pairs.
79    /// Only fields that are present and resolve to a string are included.
80    pub(crate) fn special_tokens(&self) -> Vec<(String, String)> {
81        let mut out = Vec::new();
82        for (name, field) in [
83            ("bos_token", &self.bos_token),
84            ("eos_token", &self.eos_token),
85            ("pad_token", &self.pad_token),
86            ("unk_token", &self.unk_token),
87        ] {
88            if let Some(s) = field.as_ref().and_then(TokenField::as_str) {
89                out.push((name.to_string(), s.to_string()));
90            }
91        }
92        out
93    }
94}
95
96/// Resolve the Jinja source to use from a `chat_template` field, honoring an optionally
97/// requested name (§8 rules):
98/// - explicit `name` requested → that template, else [`Error::Config`];
99/// - else a template named `"default"`;
100/// - else, if exactly one exists, that one;
101/// - else ambiguous → [`Error::Config`] listing the available names.
102pub(crate) fn resolve_template<'a>(
103    field: &'a ChatTemplateField,
104    name: Option<&str>,
105) -> Result<&'a str, Error> {
106    match field {
107        // A single string ignores any requested name — there is nothing to disambiguate.
108        ChatTemplateField::Single(s) => Ok(s),
109        ChatTemplateField::Named(list) => {
110            if list.is_empty() {
111                return Err(Error::Config("chat_template list is empty".into()));
112            }
113            if let Some(want) = name {
114                list.iter()
115                    .find(|t| t.name == want)
116                    .map(|t| t.template.as_str())
117                    .ok_or_else(|| {
118                        Error::Config(format!(
119                            "no chat_template named '{want}'; available: {}",
120                            available_names(list)
121                        ))
122                    })
123            } else if let Some(t) = list.iter().find(|t| t.name == "default") {
124                Ok(&t.template)
125            } else if list.len() == 1 {
126                Ok(&list[0].template)
127            } else {
128                Err(Error::Config(format!(
129                    "ambiguous chat_template: specify a name via builder. available: {}",
130                    available_names(list)
131                )))
132            }
133        }
134    }
135}
136
137fn available_names(list: &[NamedTemplate]) -> String {
138    list.iter()
139        .map(|t| t.name.as_str())
140        .collect::<Vec<_>>()
141        .join(", ")
142}