1pub mod basetenkenizer;
5pub mod cache;
6pub mod fastokens;
7pub mod hf;
8pub mod tiktoken;
9
10use std::hash::{DefaultHasher, Hash, Hasher};
15use std::sync::Arc;
16use std::{fs::File, io::BufReader, ops::Deref, path::Path};
17
18use anyhow::Context as _;
19pub use anyhow::{Error, Result};
20
21pub use basetenkenizer::BasetenTokenizer;
22pub use cache::{CacheTokenUsage, CacheTokenUsageFn, CachedTokenizer, L1CacheStats};
23pub use fastokens::FastTokenizer;
24pub use hf::HuggingFaceTokenizer;
25pub use tiktoken::TikTokenTokenizer;
26pub use traits::DecodeResult;
27
28pub type TokenIdType = u32;
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub struct EncodeSegment<'a> {
33 pub text: &'a str,
34 pub allow_special: bool,
39}
40
41impl<'a> EncodeSegment<'a> {
42 pub const fn new(text: &'a str, allow_special: bool) -> Self {
43 Self {
44 text,
45 allow_special,
46 }
47 }
48
49 pub const fn ordinary(text: &'a str) -> Self {
50 Self::new(text, false)
51 }
52
53 pub const fn control(text: &'a str) -> Self {
54 Self::new(text, true)
55 }
56}
57
58#[derive(Debug)]
60pub enum TokenizerType {
61 HuggingFace(String),
62 TikToken(String),
63}
64
65pub type Offsets = (usize, usize);
67
68#[derive(Debug, Clone)]
70pub enum Encoding {
71 Hf(Box<tokenizers::tokenizer::Encoding>),
73 Sp(Vec<TokenIdType>),
75}
76
77impl Encoding {
78 pub fn token_ids(&self) -> &[u32] {
79 match self {
80 Encoding::Hf(inner) => inner.get_ids(),
81 Encoding::Sp(inner) => inner,
82 }
83 }
84}
85
86impl Hash for Encoding {
87 fn hash<H: Hasher>(&self, state: &mut H) {
88 self.token_ids().hash(state);
89 }
90}
91
92pub mod traits {
93 use super::*;
94
95 pub trait Encoder: Send + Sync {
96 fn encode(&self, input: &str) -> Result<Encoding>;
97 fn encode_batch(&self, inputs: &[&str]) -> Result<Vec<Encoding>>;
98
99 fn encode_segments(&self, _segments: &[EncodeSegment<'_>]) -> Result<Encoding> {
117 Err(Error::msg(
118 "tokenizer backend does not support segmented encoding",
119 ))
120 }
121 }
122
123 #[derive(Debug, Clone, PartialEq, Eq, strum::EnumIs)]
130 pub enum DecodeResult {
131 Complete(String),
135 Partial(String),
138 }
139
140 impl DecodeResult {
141 pub fn as_str(&self) -> &str {
143 match self {
144 DecodeResult::Complete(s) | DecodeResult::Partial(s) => s,
145 }
146 }
147
148 pub fn from_decoded(text: String) -> Self {
150 if text.ends_with('\u{FFFD}') {
151 DecodeResult::Partial(text)
152 } else {
153 DecodeResult::Complete(text)
154 }
155 }
156 }
157
158 impl From<String> for DecodeResult {
159 fn from(text: String) -> Self {
160 DecodeResult::from_decoded(text)
161 }
162 }
163
164 impl From<DecodeResult> for String {
165 fn from(result: DecodeResult) -> Self {
166 match result {
167 DecodeResult::Complete(s) | DecodeResult::Partial(s) => s,
168 }
169 }
170 }
171
172 pub trait Decoder: Send + Sync {
178 fn decode(
179 &self,
180 token_ids: &[TokenIdType],
181 skip_special_tokens: bool,
182 ) -> Result<DecodeResult>;
183 }
184
185 pub trait Tokenizer: Encoder + Decoder {
186 fn validate_prefix_cache(&self) -> Result<()> {
191 Err(Error::msg("tokenizer does not support prefix caching"))
192 }
193
194 fn with_options(self, options: TokenizerOptions) -> Self
199 where
200 Self: Sized,
201 {
202 let _ = options;
203 self
204 }
205 }
208}
209
210pub fn file_json_field<T: serde::de::DeserializeOwned>(
211 json_file_path: &Path,
212 field_name: &str,
213) -> anyhow::Result<T> {
214 let file = File::open(json_file_path)
215 .with_context(|| format!("Failed to open file: {:?}", json_file_path))?;
216 let reader = BufReader::new(file);
217
218 let json_data: serde_json::Value = serde_json::from_reader(reader)
219 .with_context(|| format!("Failed to parse JSON from file: {:?}", json_file_path))?;
220
221 let map = json_data.as_object().ok_or_else(|| {
222 anyhow::anyhow!("JSON root is not an object in file: {:?}", json_file_path)
223 })?;
224
225 let field_value = map.get(field_name).ok_or_else(|| {
226 anyhow::anyhow!(
227 "Field '{}' not found in JSON file: {:?}",
228 field_name,
229 json_file_path
230 )
231 })?;
232
233 serde_json::from_value(field_value.clone()).with_context(|| {
234 format!(
235 "Failed to deserialize field '{}' (value: {:?}) to the expected type from file: {:?}",
236 field_name, field_value, json_file_path
237 )
238 })
239}
240
241pub fn log_json_err(filename: &str, json: &str, err: &serde_json::Error) {
242 const ERROR_PREFIX: &str = ">> ";
243
244 if !(err.is_syntax() || err.is_data()) {
245 return;
246 }
247
248 let line = err.line().saturating_sub(1);
249 let column = err.column().saturating_sub(1);
250
251 let json_lines: Vec<&str> = json.lines().collect();
252 if json_lines.is_empty() {
253 tracing::error!("JSON parsing error in {filename}: File is empty.");
254 return;
255 }
256
257 let start_index = line.saturating_sub(2);
258 let end_index = line.saturating_add(3).min(json_lines.len());
259
260 let mut context_lines: Vec<String> = (start_index..end_index)
261 .map(|i| {
262 if i == line {
263 format!("{ERROR_PREFIX}{}", json_lines[i])
264 } else {
265 format!("{:06} {}", i + 1, json_lines[i])
266 }
267 })
268 .collect();
269
270 let col_indicator = "_".to_string().repeat(column + ERROR_PREFIX.len()) + "^";
271 let error_in_context_idx = line - start_index;
272 if error_in_context_idx < context_lines.len() {
273 context_lines.insert(error_in_context_idx + 1, col_indicator);
274 }
275
276 tracing::error!(
277 "JSON parsing error in {filename}: Line {}, column {}:\n{}",
278 err.line(),
279 err.column(),
280 context_lines.join("\n")
281 );
282}
283
284impl Encoding {
285 pub fn get_hash(&self) -> u64 {
286 let mut hasher = DefaultHasher::new();
287 self.hash(&mut hasher);
288 hasher.finish()
289 }
290}
291
292#[derive(Debug, Clone, Copy, Default)]
296pub struct TokenizerOptions {
297 pub add_special_tokens: bool,
304}
305
306#[derive(Clone)]
308pub struct Tokenizer(Arc<dyn traits::Tokenizer>);
309
310impl Tokenizer {
311 pub fn from_file(file_path: &str) -> Result<Tokenizer> {
312 Ok(Tokenizer(create_tokenizer_from_file(file_path)?))
313 }
314
315 pub fn from_file_with_options(file_path: &str, options: TokenizerOptions) -> Result<Tokenizer> {
316 Ok(Tokenizer(create_tokenizer_from_file_with_options(
317 file_path, options,
318 )?))
319 }
320
321 pub fn decode_stream(
323 &self,
324 prompt_token_ids: &[TokenIdType],
325 skip_special_tokens: bool,
326 ) -> DecodeStream {
327 DecodeStream::new(self.0.clone(), prompt_token_ids, skip_special_tokens)
328 }
329}
330
331impl Deref for Tokenizer {
332 type Target = Arc<dyn traits::Tokenizer>;
333
334 fn deref(&self) -> &Self::Target {
335 &self.0
336 }
337}
338
339impl From<Arc<dyn traits::Tokenizer>> for Tokenizer {
340 fn from(tokenizer: Arc<dyn traits::Tokenizer>) -> Self {
341 Tokenizer(tokenizer)
342 }
343}
344
345impl<T> From<Arc<T>> for Tokenizer
346where
347 T: traits::Tokenizer + 'static, {
349 fn from(tokenizer: Arc<T>) -> Self {
350 Tokenizer(tokenizer)
351 }
352}
353
354pub fn create_tokenizer_from_file(file_path: &str) -> Result<Arc<dyn traits::Tokenizer>> {
361 create_tokenizer_from_file_with_options(file_path, Default::default())
362}
363
364pub fn create_tokenizer_from_file_with_options(
371 file_path: &str,
372 options: TokenizerOptions,
373) -> Result<Arc<dyn traits::Tokenizer>> {
374 use traits::Tokenizer as _;
375
376 let path = Path::new(file_path);
377 let extension = path
378 .extension()
379 .and_then(std::ffi::OsStr::to_str)
380 .ok_or_else(|| Error::msg("Failed to read file extension".to_string()))?;
381
382 match extension {
383 "json" => {
384 let tokenizer = HuggingFaceTokenizer::from_file(file_path)?.with_options(options);
385 Ok(Arc::new(tokenizer))
386 }
387 "model" | "tiktoken" => {
388 let tokenizer = TikTokenTokenizer::from_file_auto(file_path)?.with_options(options);
389 Ok(Arc::new(tokenizer))
390 }
391 _ => Err(Error::msg(format!(
392 "Unsupported tokenizer file type: .{extension}"
393 ))),
394 }
395}
396
397const INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET: usize = 5;
403
404pub struct DecodeStream {
410 tokenizer: Arc<dyn traits::Tokenizer>,
412
413 skip_special_tokens: bool,
414 all_token_ids: Vec<u32>,
426
427 prefix_offset: usize,
428
429 read_offset: usize,
430}
431
432impl DecodeStream {
433 pub fn new(
434 tokenizer: Arc<dyn traits::Tokenizer>,
435 prompt_token_ids: &[TokenIdType],
436 skip_special_tokens: bool,
437 ) -> Self {
438 let num_input_tokens = prompt_token_ids.len();
439 let prompt_token_ids = prompt_token_ids.to_vec();
440 Self {
441 tokenizer,
442 skip_special_tokens,
443 all_token_ids: prompt_token_ids,
444 prefix_offset: num_input_tokens
445 .saturating_sub(INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET),
446 read_offset: num_input_tokens,
447 }
448 }
449
450 pub fn step(&mut self, id: u32) -> Result<Option<String>> {
460 self.all_token_ids.push(id);
461
462 let prefix_text: String = self
463 .tokenizer
464 .decode(
465 &self.all_token_ids[self.prefix_offset..self.read_offset],
466 self.skip_special_tokens,
467 )?
468 .into();
469
470 let new_result = self.tokenizer.decode(
471 &self.all_token_ids[self.prefix_offset..],
472 self.skip_special_tokens,
473 )?;
474
475 let new_text = new_result.as_str();
476 if new_text.len() > prefix_text.len() && !new_result.is_partial() {
477 let emitted = new_text[prefix_text.len()..].to_string();
478
479 self.prefix_offset = self.read_offset;
480 self.read_offset = self.all_token_ids.len();
481
482 Ok(Some(emitted))
483 } else {
484 Ok(None)
485 }
486 }
487}
488
489pub struct Sequence {
491 tokenizer: Tokenizer,
493
494 token_ids: Vec<TokenIdType>,
496
497 prefix_offset: usize,
499
500 read_offset: usize,
502}
503
504impl std::fmt::Debug for Sequence {
505 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
506 f.debug_struct("Sequence")
507 .field("tokenizer", &"Arc<dyn Tokenizer>")
508 .field(
509 "token_ids",
510 &format_args!("{}", {
511 let token_ids = self.token_ids();
512 if token_ids.len() <= 20 {
513 format!("{:?}", token_ids)
514 } else {
515 let first_ten = &token_ids[..10];
516 let last_ten = &token_ids[token_ids.len() - 10..];
517 format!("{:?} ... {:?}", first_ten, last_ten)
518 }
519 }),
520 )
521 .field("prefix_offset", &self.prefix_offset)
522 .field("read_offset", &self.read_offset)
523 .field("token count", &self.token_ids.len())
524 .finish()
525 }
526}
527
528impl Sequence {
529 pub fn new(tokenizer: Tokenizer) -> Self {
530 Self {
531 tokenizer,
532 token_ids: Vec::new(),
533 prefix_offset: 0,
534 read_offset: 0,
535 }
536 }
537
538 pub fn is_empty(&self) -> bool {
539 self.token_ids.is_empty()
540 }
541
542 pub fn len(&self) -> usize {
543 self.token_ids.len()
544 }
545
546 pub fn clear(&mut self) {
547 self.token_ids.clear();
548 self.prefix_offset = 0;
549 self.read_offset = 0;
550 }
551
552 pub fn append_text(&mut self, input: &str) -> Result<()> {
553 let encoding = self.tokenizer.encode(input)?;
558 self.token_ids.extend(encoding.token_ids());
559 Ok(())
560 }
561
562 pub fn append_token_id(&mut self, token_id: TokenIdType) -> Result<String> {
566 self.token_ids.push(token_id);
567 let prefix_text: String = self
570 .tokenizer
571 .decode(&self.token_ids[self.prefix_offset..self.read_offset], false)?
572 .into();
573
574 let new_result = self
575 .tokenizer
576 .decode(&self.token_ids[self.prefix_offset..], false)?;
577
578 let new_text = new_result.as_str();
579
580 let mut prefix_text_len = prefix_text.len();
584 while !new_text.is_char_boundary(prefix_text_len) && prefix_text_len > 0 {
585 prefix_text_len -= 1;
586 }
587 let prefix_text_len = prefix_text_len;
588
589 if new_text.len() > prefix_text.len() {
590 if new_result.is_partial() {
591 return Ok("".to_string());
592 } else {
593 let new_text = new_text[prefix_text_len..]
595 .to_string()
596 .replace('\u{FFFD}', "");
597 self.prefix_offset = self.read_offset;
598 self.read_offset = self.token_ids.len();
599 return Ok(new_text);
600 }
601 }
602
603 Ok("".to_string())
604 }
605
606 pub fn tokenizer(&self) -> Tokenizer {
607 self.tokenizer.clone()
608 }
609
610 pub fn token_ids(&self) -> &[TokenIdType] {
611 &self.token_ids
612 }
613
614 pub fn text(&self) -> Result<String> {
615 Ok(self.tokenizer.decode(&self.token_ids, false)?.into())
619 }
620}
621
622pub enum SequenceDecoderOutput {
625 Text(String),
627
628 Held,
631
632 Stopped,
635
636 StoppedWithText(String),
640}
641
642#[derive(Debug)]
648pub struct StopSequenceDecoder {
649 sequence: Sequence,
651
652 stop_token_ids_visible: Vec<TokenIdType>,
655
656 stop_token_ids_hidden: Vec<TokenIdType>,
659
660 #[allow(dead_code)]
663 stop_sequences_visible: Vec<String>,
664
665 stop_sequences_hidden: Vec<String>,
668
669 stopped: bool,
672
673 state: String,
676}
677
678impl StopSequenceDecoder {
679 pub fn builder(tokenizer: Tokenizer) -> StopSequenceDecoderBuilder {
681 StopSequenceDecoderBuilder::new(tokenizer)
682 }
683
684 pub fn append_token_id(&mut self, token_id: TokenIdType) -> Result<SequenceDecoderOutput> {
686 if self.stopped {
687 return Err(Error::msg("Decoder is stopped"));
688 }
689
690 let text = self.sequence.append_token_id(token_id)?;
692
693 self.state.push_str(text.as_str());
695
696 let mut stop: bool = false;
697 let mut visible: bool = false;
698
699 if self.stop_token_ids_visible.contains(&token_id) {
700 stop = true;
701 visible = true;
702 }
703
704 if self.stop_token_ids_hidden.contains(&token_id) {
705 stop = true;
706 visible = false;
707 }
708
709 if stop {
710 self.stopped = true;
711 let state = std::mem::take(&mut self.state);
712 if visible {
713 return Ok(SequenceDecoderOutput::StoppedWithText(state));
714 }
715 return Ok(SequenceDecoderOutput::Stopped);
716 }
717
718 for stop_sequence in self.stop_sequences_hidden.iter() {
720 if stop_sequence.starts_with(&self.state) {
721 if stop_sequence == &self.state {
722 self.stopped = true;
724 return Ok(SequenceDecoderOutput::Stopped);
725 } else {
726 return Ok(SequenceDecoderOutput::Held);
727 }
728 }
729 }
730
731 let state = std::mem::take(&mut self.state);
732 Ok(SequenceDecoderOutput::Text(state))
733 }
734
735 pub fn is_empty(&self) -> bool {
736 self.sequence.token_ids.is_empty()
737 }
738
739 pub fn len(&self) -> usize {
740 self.sequence.token_ids.len()
741 }
742
743 pub fn is_complete(&self) -> bool {
744 self.stopped
745 }
746
747 pub fn close(&mut self) {
748 self.stopped = true;
749 }
750}
751
752pub struct StopSequenceDecoderBuilder {
753 tokenizer: Tokenizer,
754 stop_token_ids_visible: Vec<TokenIdType>,
755 stop_token_ids_hidden: Vec<TokenIdType>,
756 stop_sequences_visible: Vec<String>,
757 stop_sequences_hidden: Vec<String>,
758}
759
760impl StopSequenceDecoderBuilder {
761 pub fn new(tokenizer: Tokenizer) -> Self {
762 Self {
763 tokenizer,
764 stop_token_ids_visible: Vec::new(),
765 stop_token_ids_hidden: Vec::new(),
766 stop_sequences_visible: Vec::new(),
767 stop_sequences_hidden: Vec::new(),
768 }
769 }
770
771 pub fn add_stop_token_id_visible(mut self, token_id: TokenIdType) -> Self {
773 self.stop_token_ids_visible.push(token_id);
774 self
775 }
776
777 pub fn add_stop_token_ids_visible(mut self, token_ids: &[TokenIdType]) -> Self {
780 self.stop_token_ids_visible.extend(token_ids);
781 self
782 }
783
784 pub fn add_stop_token_id_hidden(mut self, token_id: TokenIdType) -> Self {
786 self.stop_token_ids_hidden.push(token_id);
787 self
788 }
789
790 pub fn add_stop_token_ids_hidden(mut self, token_ids: &[TokenIdType]) -> Self {
793 self.stop_token_ids_hidden.extend(token_ids);
794 self
795 }
796
797 pub fn add_stop_sequence_visible(mut self, text: &str) -> Self {
798 self.stop_sequences_visible.push(text.to_string());
799 self
800 }
801
802 pub fn add_stop_sequences_visible(mut self, strings: &[&str]) -> Self {
803 self.stop_sequences_visible
804 .extend(strings.iter().map(|text| text.to_string()));
805 self
806 }
807
808 pub fn add_stop_sequence_hidden(mut self, text: &str) -> Self {
809 self.stop_sequences_hidden.push(text.to_string());
810 self
811 }
812
813 pub fn add_stop_sequences_hidden(mut self, strings: &[&str]) -> Self {
814 self.stop_sequences_hidden
815 .extend(strings.iter().map(|text| text.to_string()));
816 self
817 }
818
819 pub fn build(self) -> Result<StopSequenceDecoder> {
820 Ok(StopSequenceDecoder {
821 sequence: Sequence::new(self.tokenizer.clone()),
822 stop_token_ids_visible: self.stop_token_ids_visible,
823 stop_token_ids_hidden: self.stop_token_ids_hidden,
824 stop_sequences_visible: self.stop_sequences_visible,
825 stop_sequences_hidden: self.stop_sequences_hidden,
826 stopped: false,
827 state: String::new(),
828 })
829 }
830}