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
50#[derive(Debug)]
52pub enum TokenizerType {
53 HuggingFace(String),
54 TikToken(String),
55}
56
57pub type Offsets = (usize, usize);
59
60#[derive(Debug, Clone)]
62pub enum Encoding {
63 Hf(Box<tokenizers::tokenizer::Encoding>),
65 Sp(Vec<TokenIdType>),
67}
68
69impl Encoding {
70 pub fn token_ids(&self) -> &[u32] {
71 match self {
72 Encoding::Hf(inner) => inner.get_ids(),
73 Encoding::Sp(inner) => inner,
74 }
75 }
76}
77
78impl Hash for Encoding {
79 fn hash<H: Hasher>(&self, state: &mut H) {
80 self.token_ids().hash(state);
81 }
82}
83
84pub mod traits {
85 use super::*;
86
87 pub trait Encoder: Send + Sync {
88 fn encode(&self, input: &str) -> Result<Encoding>;
89 fn encode_batch(&self, inputs: &[&str]) -> Result<Vec<Encoding>>;
90
91 fn encode_segments(&self, _segments: &[EncodeSegment<'_>]) -> Result<Encoding> {
109 Err(Error::msg(
110 "tokenizer backend does not support segmented encoding",
111 ))
112 }
113 }
114
115 #[derive(Debug, Clone, PartialEq, Eq, strum::EnumIs)]
122 pub enum DecodeResult {
123 Complete(String),
127 Partial(String),
130 }
131
132 impl DecodeResult {
133 pub fn as_str(&self) -> &str {
135 match self {
136 DecodeResult::Complete(s) | DecodeResult::Partial(s) => s,
137 }
138 }
139
140 pub fn from_decoded(text: String) -> Self {
142 if text.ends_with('\u{FFFD}') {
143 DecodeResult::Partial(text)
144 } else {
145 DecodeResult::Complete(text)
146 }
147 }
148 }
149
150 impl From<String> for DecodeResult {
151 fn from(text: String) -> Self {
152 DecodeResult::from_decoded(text)
153 }
154 }
155
156 impl From<DecodeResult> for String {
157 fn from(result: DecodeResult) -> Self {
158 match result {
159 DecodeResult::Complete(s) | DecodeResult::Partial(s) => s,
160 }
161 }
162 }
163
164 pub trait Decoder: Send + Sync {
170 fn decode(
171 &self,
172 token_ids: &[TokenIdType],
173 skip_special_tokens: bool,
174 ) -> Result<DecodeResult>;
175 }
176
177 pub trait Tokenizer: Encoder + Decoder {
178 fn validate_prefix_cache(&self) -> Result<()> {
183 Err(Error::msg("tokenizer does not support prefix caching"))
184 }
185
186 fn with_options(self, options: TokenizerOptions) -> Self
191 where
192 Self: Sized,
193 {
194 let _ = options;
195 self
196 }
197 }
200}
201
202pub fn file_json_field<T: serde::de::DeserializeOwned>(
203 json_file_path: &Path,
204 field_name: &str,
205) -> anyhow::Result<T> {
206 let file = File::open(json_file_path)
207 .with_context(|| format!("Failed to open file: {:?}", json_file_path))?;
208 let reader = BufReader::new(file);
209
210 let json_data: serde_json::Value = serde_json::from_reader(reader)
211 .with_context(|| format!("Failed to parse JSON from file: {:?}", json_file_path))?;
212
213 let map = json_data.as_object().ok_or_else(|| {
214 anyhow::anyhow!("JSON root is not an object in file: {:?}", json_file_path)
215 })?;
216
217 let field_value = map.get(field_name).ok_or_else(|| {
218 anyhow::anyhow!(
219 "Field '{}' not found in JSON file: {:?}",
220 field_name,
221 json_file_path
222 )
223 })?;
224
225 serde_json::from_value(field_value.clone()).with_context(|| {
226 format!(
227 "Failed to deserialize field '{}' (value: {:?}) to the expected type from file: {:?}",
228 field_name, field_value, json_file_path
229 )
230 })
231}
232
233pub fn log_json_err(filename: &str, json: &str, err: &serde_json::Error) {
234 const ERROR_PREFIX: &str = ">> ";
235
236 if !(err.is_syntax() || err.is_data()) {
237 return;
238 }
239
240 let line = err.line().saturating_sub(1);
241 let column = err.column().saturating_sub(1);
242
243 let json_lines: Vec<&str> = json.lines().collect();
244 if json_lines.is_empty() {
245 tracing::error!("JSON parsing error in {filename}: File is empty.");
246 return;
247 }
248
249 let start_index = line.saturating_sub(2);
250 let end_index = line.saturating_add(3).min(json_lines.len());
251
252 let mut context_lines: Vec<String> = (start_index..end_index)
253 .map(|i| {
254 if i == line {
255 format!("{ERROR_PREFIX}{}", json_lines[i])
256 } else {
257 format!("{:06} {}", i + 1, json_lines[i])
258 }
259 })
260 .collect();
261
262 let col_indicator = "_".to_string().repeat(column + ERROR_PREFIX.len()) + "^";
263 let error_in_context_idx = line - start_index;
264 if error_in_context_idx < context_lines.len() {
265 context_lines.insert(error_in_context_idx + 1, col_indicator);
266 }
267
268 tracing::error!(
269 "JSON parsing error in {filename}: Line {}, column {}:\n{}",
270 err.line(),
271 err.column(),
272 context_lines.join("\n")
273 );
274}
275
276impl Encoding {
277 pub fn get_hash(&self) -> u64 {
278 let mut hasher = DefaultHasher::new();
279 self.hash(&mut hasher);
280 hasher.finish()
281 }
282}
283
284#[derive(Debug, Clone, Copy, Default)]
288pub struct TokenizerOptions {
289 pub add_special_tokens: bool,
296}
297
298#[derive(Clone)]
300pub struct Tokenizer(Arc<dyn traits::Tokenizer>);
301
302impl Tokenizer {
303 pub fn from_file(file_path: &str) -> Result<Tokenizer> {
304 Ok(Tokenizer(create_tokenizer_from_file(file_path)?))
305 }
306
307 pub fn from_file_with_options(file_path: &str, options: TokenizerOptions) -> Result<Tokenizer> {
308 Ok(Tokenizer(create_tokenizer_from_file_with_options(
309 file_path, options,
310 )?))
311 }
312
313 pub fn decode_stream(
315 &self,
316 prompt_token_ids: &[TokenIdType],
317 skip_special_tokens: bool,
318 ) -> DecodeStream {
319 DecodeStream::new(self.0.clone(), prompt_token_ids, skip_special_tokens)
320 }
321}
322
323impl Deref for Tokenizer {
324 type Target = Arc<dyn traits::Tokenizer>;
325
326 fn deref(&self) -> &Self::Target {
327 &self.0
328 }
329}
330
331impl From<Arc<dyn traits::Tokenizer>> for Tokenizer {
332 fn from(tokenizer: Arc<dyn traits::Tokenizer>) -> Self {
333 Tokenizer(tokenizer)
334 }
335}
336
337impl<T> From<Arc<T>> for Tokenizer
338where
339 T: traits::Tokenizer + 'static, {
341 fn from(tokenizer: Arc<T>) -> Self {
342 Tokenizer(tokenizer)
343 }
344}
345
346pub fn create_tokenizer_from_file(file_path: &str) -> Result<Arc<dyn traits::Tokenizer>> {
353 create_tokenizer_from_file_with_options(file_path, Default::default())
354}
355
356pub fn create_tokenizer_from_file_with_options(
363 file_path: &str,
364 options: TokenizerOptions,
365) -> Result<Arc<dyn traits::Tokenizer>> {
366 use traits::Tokenizer as _;
367
368 let path = Path::new(file_path);
369 let extension = path
370 .extension()
371 .and_then(std::ffi::OsStr::to_str)
372 .ok_or_else(|| Error::msg("Failed to read file extension".to_string()))?;
373
374 match extension {
375 "json" => {
376 let tokenizer = HuggingFaceTokenizer::from_file(file_path)?.with_options(options);
377 Ok(Arc::new(tokenizer))
378 }
379 "model" | "tiktoken" => {
380 let tokenizer = TikTokenTokenizer::from_file_auto(file_path)?.with_options(options);
381 Ok(Arc::new(tokenizer))
382 }
383 _ => Err(Error::msg(format!(
384 "Unsupported tokenizer file type: .{extension}"
385 ))),
386 }
387}
388
389const INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET: usize = 5;
395
396pub struct DecodeStream {
402 tokenizer: Arc<dyn traits::Tokenizer>,
404
405 skip_special_tokens: bool,
406 all_token_ids: Vec<u32>,
418
419 prefix_offset: usize,
420
421 read_offset: usize,
422}
423
424impl DecodeStream {
425 pub fn new(
426 tokenizer: Arc<dyn traits::Tokenizer>,
427 prompt_token_ids: &[TokenIdType],
428 skip_special_tokens: bool,
429 ) -> Self {
430 let num_input_tokens = prompt_token_ids.len();
431 let prompt_token_ids = prompt_token_ids.to_vec();
432 Self {
433 tokenizer,
434 skip_special_tokens,
435 all_token_ids: prompt_token_ids,
436 prefix_offset: num_input_tokens
437 .saturating_sub(INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET),
438 read_offset: num_input_tokens,
439 }
440 }
441
442 pub fn step(&mut self, id: u32) -> Result<Option<String>> {
452 self.all_token_ids.push(id);
453
454 let prefix_text: String = self
455 .tokenizer
456 .decode(
457 &self.all_token_ids[self.prefix_offset..self.read_offset],
458 self.skip_special_tokens,
459 )?
460 .into();
461
462 let new_result = self.tokenizer.decode(
463 &self.all_token_ids[self.prefix_offset..],
464 self.skip_special_tokens,
465 )?;
466
467 let new_text = new_result.as_str();
468 if new_text.len() > prefix_text.len() && !new_result.is_partial() {
469 let emitted = new_text[prefix_text.len()..].to_string();
470
471 self.prefix_offset = self.read_offset;
472 self.read_offset = self.all_token_ids.len();
473
474 Ok(Some(emitted))
475 } else {
476 Ok(None)
477 }
478 }
479}
480
481pub struct Sequence {
483 tokenizer: Tokenizer,
485
486 token_ids: Vec<TokenIdType>,
488
489 prefix_offset: usize,
491
492 read_offset: usize,
494}
495
496impl std::fmt::Debug for Sequence {
497 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
498 f.debug_struct("Sequence")
499 .field("tokenizer", &"Arc<dyn Tokenizer>")
500 .field(
501 "token_ids",
502 &format_args!("{}", {
503 let token_ids = self.token_ids();
504 if token_ids.len() <= 20 {
505 format!("{:?}", token_ids)
506 } else {
507 let first_ten = &token_ids[..10];
508 let last_ten = &token_ids[token_ids.len() - 10..];
509 format!("{:?} ... {:?}", first_ten, last_ten)
510 }
511 }),
512 )
513 .field("prefix_offset", &self.prefix_offset)
514 .field("read_offset", &self.read_offset)
515 .field("token count", &self.token_ids.len())
516 .finish()
517 }
518}
519
520impl Sequence {
521 pub fn new(tokenizer: Tokenizer) -> Self {
522 Self {
523 tokenizer,
524 token_ids: Vec::new(),
525 prefix_offset: 0,
526 read_offset: 0,
527 }
528 }
529
530 pub fn is_empty(&self) -> bool {
531 self.token_ids.is_empty()
532 }
533
534 pub fn len(&self) -> usize {
535 self.token_ids.len()
536 }
537
538 pub fn clear(&mut self) {
539 self.token_ids.clear();
540 self.prefix_offset = 0;
541 self.read_offset = 0;
542 }
543
544 pub fn append_text(&mut self, input: &str) -> Result<()> {
545 let encoding = self.tokenizer.encode(input)?;
550 self.token_ids.extend(encoding.token_ids());
551 Ok(())
552 }
553
554 pub fn append_token_id(&mut self, token_id: TokenIdType) -> Result<String> {
558 self.token_ids.push(token_id);
559 let prefix_text: String = self
562 .tokenizer
563 .decode(&self.token_ids[self.prefix_offset..self.read_offset], false)?
564 .into();
565
566 let new_result = self
567 .tokenizer
568 .decode(&self.token_ids[self.prefix_offset..], false)?;
569
570 let new_text = new_result.as_str();
571
572 let mut prefix_text_len = prefix_text.len();
576 while !new_text.is_char_boundary(prefix_text_len) && prefix_text_len > 0 {
577 prefix_text_len -= 1;
578 }
579 let prefix_text_len = prefix_text_len;
580
581 if new_text.len() > prefix_text.len() {
582 if new_result.is_partial() {
583 return Ok("".to_string());
584 } else {
585 let new_text = new_text[prefix_text_len..]
587 .to_string()
588 .replace('\u{FFFD}', "");
589 self.prefix_offset = self.read_offset;
590 self.read_offset = self.token_ids.len();
591 return Ok(new_text);
592 }
593 }
594
595 Ok("".to_string())
596 }
597
598 pub fn tokenizer(&self) -> Tokenizer {
599 self.tokenizer.clone()
600 }
601
602 pub fn token_ids(&self) -> &[TokenIdType] {
603 &self.token_ids
604 }
605
606 pub fn text(&self) -> Result<String> {
607 Ok(self.tokenizer.decode(&self.token_ids, false)?.into())
611 }
612}
613
614pub enum SequenceDecoderOutput {
617 Text(String),
619
620 Held,
623
624 Stopped,
627
628 StoppedWithText(String),
632}
633
634#[derive(Debug)]
640pub struct StopSequenceDecoder {
641 sequence: Sequence,
643
644 stop_token_ids_visible: Vec<TokenIdType>,
647
648 stop_token_ids_hidden: Vec<TokenIdType>,
651
652 #[allow(dead_code)]
655 stop_sequences_visible: Vec<String>,
656
657 stop_sequences_hidden: Vec<String>,
660
661 stopped: bool,
664
665 state: String,
668}
669
670impl StopSequenceDecoder {
671 pub fn builder(tokenizer: Tokenizer) -> StopSequenceDecoderBuilder {
673 StopSequenceDecoderBuilder::new(tokenizer)
674 }
675
676 pub fn append_token_id(&mut self, token_id: TokenIdType) -> Result<SequenceDecoderOutput> {
678 if self.stopped {
679 return Err(Error::msg("Decoder is stopped"));
680 }
681
682 let text = self.sequence.append_token_id(token_id)?;
684
685 self.state.push_str(text.as_str());
687
688 let mut stop: bool = false;
689 let mut visible: bool = false;
690
691 if self.stop_token_ids_visible.contains(&token_id) {
692 stop = true;
693 visible = true;
694 }
695
696 if self.stop_token_ids_hidden.contains(&token_id) {
697 stop = true;
698 visible = false;
699 }
700
701 if stop {
702 self.stopped = true;
703 let state = std::mem::take(&mut self.state);
704 if visible {
705 return Ok(SequenceDecoderOutput::StoppedWithText(state));
706 }
707 return Ok(SequenceDecoderOutput::Stopped);
708 }
709
710 for stop_sequence in self.stop_sequences_hidden.iter() {
712 if stop_sequence.starts_with(&self.state) {
713 if stop_sequence == &self.state {
714 self.stopped = true;
716 return Ok(SequenceDecoderOutput::Stopped);
717 } else {
718 return Ok(SequenceDecoderOutput::Held);
719 }
720 }
721 }
722
723 let state = std::mem::take(&mut self.state);
724 Ok(SequenceDecoderOutput::Text(state))
725 }
726
727 pub fn is_empty(&self) -> bool {
728 self.sequence.token_ids.is_empty()
729 }
730
731 pub fn len(&self) -> usize {
732 self.sequence.token_ids.len()
733 }
734
735 pub fn is_complete(&self) -> bool {
736 self.stopped
737 }
738
739 pub fn close(&mut self) {
740 self.stopped = true;
741 }
742}
743
744pub struct StopSequenceDecoderBuilder {
745 tokenizer: Tokenizer,
746 stop_token_ids_visible: Vec<TokenIdType>,
747 stop_token_ids_hidden: Vec<TokenIdType>,
748 stop_sequences_visible: Vec<String>,
749 stop_sequences_hidden: Vec<String>,
750}
751
752impl StopSequenceDecoderBuilder {
753 pub fn new(tokenizer: Tokenizer) -> Self {
754 Self {
755 tokenizer,
756 stop_token_ids_visible: Vec::new(),
757 stop_token_ids_hidden: Vec::new(),
758 stop_sequences_visible: Vec::new(),
759 stop_sequences_hidden: Vec::new(),
760 }
761 }
762
763 pub fn add_stop_token_id_visible(mut self, token_id: TokenIdType) -> Self {
765 self.stop_token_ids_visible.push(token_id);
766 self
767 }
768
769 pub fn add_stop_token_ids_visible(mut self, token_ids: &[TokenIdType]) -> Self {
772 self.stop_token_ids_visible.extend(token_ids);
773 self
774 }
775
776 pub fn add_stop_token_id_hidden(mut self, token_id: TokenIdType) -> Self {
778 self.stop_token_ids_hidden.push(token_id);
779 self
780 }
781
782 pub fn add_stop_token_ids_hidden(mut self, token_ids: &[TokenIdType]) -> Self {
785 self.stop_token_ids_hidden.extend(token_ids);
786 self
787 }
788
789 pub fn add_stop_sequence_visible(mut self, text: &str) -> Self {
790 self.stop_sequences_visible.push(text.to_string());
791 self
792 }
793
794 pub fn add_stop_sequences_visible(mut self, strings: &[&str]) -> Self {
795 self.stop_sequences_visible
796 .extend(strings.iter().map(|text| text.to_string()));
797 self
798 }
799
800 pub fn add_stop_sequence_hidden(mut self, text: &str) -> Self {
801 self.stop_sequences_hidden.push(text.to_string());
802 self
803 }
804
805 pub fn add_stop_sequences_hidden(mut self, strings: &[&str]) -> Self {
806 self.stop_sequences_hidden
807 .extend(strings.iter().map(|text| text.to_string()));
808 self
809 }
810
811 pub fn build(self) -> Result<StopSequenceDecoder> {
812 Ok(StopSequenceDecoder {
813 sequence: Sequence::new(self.tokenizer.clone()),
814 stop_token_ids_visible: self.stop_token_ids_visible,
815 stop_token_ids_hidden: self.stop_token_ids_hidden,
816 stop_sequences_visible: self.stop_sequences_visible,
817 stop_sequences_hidden: self.stop_sequences_hidden,
818 stopped: false,
819 state: String::new(),
820 })
821 }
822}