dynamo_renderer/template/
tokcfg.rs1use 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#[derive(Debug, Deserialize)]
41pub struct ChatTemplateValue(
42 #[serde(with = "either::serde_untagged")] pub Either<String, Vec<HashMap<String, String>>>,
43);
44
45#[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)]
54pub struct ChatTemplate {
56 pub bos_token: Option<BeginEndUnkTok>,
57 pub eos_token: Option<BeginEndUnkTok>,
58 pub unk_token: Option<BeginEndUnkTok>,
59
60 pub chat_template: Option<ChatTemplateValue>,
64
65 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
114struct 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
149pub 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 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}