hf_chat_template/
config.rs1use serde::{Deserialize, Serialize};
9use serde_json::{Map, Value as Json};
10
11use crate::error::Error;
12
13#[derive(Clone, Debug, Default, Serialize, Deserialize)]
15pub struct TokenizerConfig {
16 #[serde(default, skip_serializing_if = "Option::is_none")]
18 pub chat_template: Option<ChatTemplateField>,
19
20 #[serde(default, skip_serializing_if = "Option::is_none")]
22 pub bos_token: Option<TokenField>,
23 #[serde(default, skip_serializing_if = "Option::is_none")]
25 pub eos_token: Option<TokenField>,
26 #[serde(default, skip_serializing_if = "Option::is_none")]
28 pub pad_token: Option<TokenField>,
29 #[serde(default, skip_serializing_if = "Option::is_none")]
31 pub unk_token: Option<TokenField>,
32
33 #[serde(default, flatten)]
35 pub extra: Map<String, Json>,
36}
37
38#[derive(Clone, Debug, Serialize, Deserialize)]
40#[serde(untagged)]
41pub enum ChatTemplateField {
42 Single(String),
44 Named(Vec<NamedTemplate>),
46}
47
48#[derive(Clone, Debug, Serialize, Deserialize)]
50pub struct NamedTemplate {
51 pub name: String,
53 pub template: String,
55}
56
57#[derive(Clone, Debug, Serialize, Deserialize)]
59#[serde(untagged)]
60pub enum TokenField {
61 Str(String),
63 Obj(Map<String, Json>),
65}
66
67impl TokenField {
68 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 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
96pub(crate) fn resolve_template<'a>(
103 field: &'a ChatTemplateField,
104 name: Option<&str>,
105) -> Result<&'a str, Error> {
106 match field {
107 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}