1use ahash::AHashMap;
13use std::{
14 fs::{read_to_string, File},
15 io::{prelude::*, BufReader},
16 ops::{Deref, DerefMut},
17 path::{Path, PathBuf},
18};
19
20use serde::de::DeserializeOwned;
21use serde::{Deserialize, Serialize};
22
23use crate::utils::iter::ResultShunt;
24use crate::utils::parallelism::*;
25use crate::utils::progress::{ProgressBar, ProgressStyle};
26
27mod added_vocabulary;
28mod encoding;
29pub mod normalizer;
30pub mod pattern;
31pub mod pre_tokenizer;
32mod serialization;
33
34pub use crate::decoders::DecoderWrapper;
36pub use crate::models::ModelWrapper;
37pub use crate::normalizers::NormalizerWrapper;
38pub use crate::pre_tokenizers::PreTokenizerWrapper;
39pub use crate::processors::PostProcessorWrapper;
40pub use crate::utils::iter::LinesWithEnding;
42pub use crate::utils::padding::{pad_encodings, PaddingDirection, PaddingParams, PaddingStrategy};
43pub use crate::utils::truncation::{
44 truncate_encodings, TruncationDirection, TruncationParams, TruncationStrategy,
45};
46pub use added_vocabulary::*;
47pub use encoding::*;
48pub use normalizer::{NormalizedString, OffsetReferential, SplitDelimiterBehavior};
49pub use pre_tokenizer::*;
50
51pub type Error = Box<dyn std::error::Error + Send + Sync>;
52pub type Result<T> = std::result::Result<T, Error>;
53pub type Offsets = (usize, usize);
54
55pub trait Normalizer {
57 fn normalize(&self, normalized: &mut NormalizedString) -> Result<()>;
58}
59
60pub trait PreTokenizer {
66 fn pre_tokenize(&self, pretokenized: &mut PreTokenizedString) -> Result<()>;
67}
68
69pub trait Model {
71 type Trainer: Trainer + Sync;
72 fn tokenize(&self, sequence: &str) -> Result<Vec<Token>>;
75 fn token_to_id(&self, token: &str) -> Option<u32>;
77 fn id_to_token(&self, id: u32) -> Option<String>;
79 fn get_vocab(&self) -> HashMap<String, u32>;
81 fn get_vocab_size(&self) -> usize;
83 fn save(&self, folder: &Path, prefix: Option<&str>) -> Result<Vec<PathBuf>>;
86 fn get_trainer(&self) -> <Self as Model>::Trainer;
88}
89
90pub trait PostProcessor {
93 fn added_tokens(&self, is_pair: bool) -> usize;
95 fn process(
97 &self,
98 encoding: Encoding,
99 pair_encoding: Option<Encoding>,
100 add_special_tokens: bool,
101 ) -> Result<Encoding> {
102 let mut encodings = if let Some(pair_encoding) = pair_encoding {
103 vec![encoding, pair_encoding]
104 } else {
105 vec![encoding]
106 };
107 encodings.iter_mut().enumerate().for_each(|(i, encoding)| {
108 encoding.set_sequence_id(i);
109 encoding
110 .get_overflowing_mut()
111 .iter_mut()
112 .for_each(|encoding| encoding.set_sequence_id(i));
113 encoding.set_type_ids(vec![i as u32; encoding.len()]);
114 });
115
116 let encodings = self.process_encodings(encodings, add_special_tokens)?;
117 Ok(Encoding::merge(encodings, false))
118 }
119
120 fn process_encodings(
122 &self,
123 encodings: Vec<Encoding>,
124 add_special_tokens: bool,
125 ) -> Result<Vec<Encoding>>;
126}
127impl dyn PostProcessor {
128 pub fn default_process(
129 encodings: Vec<Encoding>,
130 _add_special_tokens: bool,
131 ) -> Result<Vec<Encoding>> {
132 match encodings.len() {
133 1 => Ok(encodings),
134 _ => {
135 let mut final_encoding = Encoding::default();
136 for (i, mut encoding) in encodings.into_iter().enumerate() {
137 encoding.set_sequence_id(i);
138 final_encoding.merge_with(encoding, false);
139 }
140 Ok(vec![final_encoding])
141 }
142 }
143 }
144}
145
146#[derive(thiserror::Error, Debug)]
147pub enum ProcessorError {
148 #[error("encodings vector length must be either 1 or 2")]
149 InvalidEncodingsVecLength,
150}
151
152pub trait Decoder {
154 fn decode(&self, tokens: Vec<String>) -> Result<String> {
155 let results = self.decode_chain(tokens)?;
156 Ok(results.join(""))
157 }
158 fn decode_chain(&self, tokens: Vec<String>) -> Result<Vec<String>>;
159}
160
161pub trait Trainer {
164 type Model: Model + Sized;
165 fn should_show_progress(&self) -> bool;
167 fn train(&self, model: &mut Self::Model) -> Result<Vec<AddedToken>>;
170 fn feed<I, S, F>(&mut self, iterator: I, process: F) -> Result<()>
173 where
174 I: Iterator<Item = S> + Send,
175 S: AsRef<str> + Send,
176 F: Fn(&str) -> Result<Vec<String>> + Sync;
177}
178
179#[derive(Debug, Clone, PartialEq, Eq)]
180pub struct Token {
181 pub id: u32,
182 pub value: String,
183 pub offsets: (usize, usize),
184}
185impl Token {
186 pub fn new(id: u32, value: String, offsets: (usize, usize)) -> Self {
187 Self { id, value, offsets }
188 }
189}
190
191use std::borrow::Cow;
192use std::collections::HashMap;
193
194#[derive(Debug, Clone)]
195pub enum InputSequence<'s> {
196 Raw(Cow<'s, str>),
197 PreTokenized(Cow<'s, [&'s str]>),
198 PreTokenizedOwned(Cow<'s, [String]>),
199 PreTokenizedCow(Cow<'s, [Cow<'s, str>]>),
200}
201
202impl<'s> From<Cow<'s, str>> for InputSequence<'s> {
203 fn from(input: Cow<'s, str>) -> Self {
204 Self::Raw(input)
205 }
206}
207
208impl<'s> From<&'s str> for InputSequence<'s> {
209 fn from(input: &'s str) -> Self {
210 Self::Raw(Cow::Borrowed(input))
211 }
212}
213
214impl From<String> for InputSequence<'_> {
215 fn from(input: String) -> Self {
216 Self::Raw(Cow::Owned(input))
217 }
218}
219
220impl<'s> From<&'s [&'s str]> for InputSequence<'s> {
221 fn from(input: &'s [&'s str]) -> Self {
222 Self::PreTokenized(Cow::Borrowed(input))
223 }
224}
225
226impl<'s> From<Vec<&'s str>> for InputSequence<'s> {
227 fn from(input: Vec<&'s str>) -> Self {
228 Self::PreTokenized(Cow::Owned(input))
229 }
230}
231
232impl<'s> From<&'s [String]> for InputSequence<'s> {
233 fn from(input: &'s [String]) -> Self {
234 Self::PreTokenizedOwned(Cow::Borrowed(input))
235 }
236}
237
238impl From<Vec<String>> for InputSequence<'_> {
239 fn from(input: Vec<String>) -> Self {
240 Self::PreTokenizedOwned(Cow::Owned(input))
241 }
242}
243
244impl<'s> From<Vec<Cow<'s, str>>> for InputSequence<'s> {
245 fn from(input: Vec<Cow<'s, str>>) -> Self {
246 Self::PreTokenizedCow(Cow::Owned(input))
247 }
248}
249
250impl<'s> From<&'s [Cow<'s, str>]> for InputSequence<'s> {
251 fn from(input: &'s [Cow<'s, str>]) -> Self {
252 Self::PreTokenizedCow(Cow::Borrowed(input))
253 }
254}
255
256#[derive(Debug, Clone)]
257pub enum EncodeInput<'s> {
258 Single(InputSequence<'s>),
259 Dual(InputSequence<'s>, InputSequence<'s>),
260}
261
262impl<'s, I: Into<InputSequence<'s>>> From<I> for EncodeInput<'s> {
263 fn from(input: I) -> Self {
264 Self::Single(input.into())
265 }
266}
267
268impl<'s, I1, I2> From<(I1, I2)> for EncodeInput<'s>
269where
270 I1: Into<InputSequence<'s>>,
271 I2: Into<InputSequence<'s>>,
272{
273 fn from(input: (I1, I2)) -> Self {
274 Self::Dual(input.0.into(), input.1.into())
275 }
276}
277
278#[derive(thiserror::Error, Debug)]
279#[error("{0}")]
280pub struct BuilderError(String);
281
282pub struct TokenizerBuilder<M, N, PT, PP, D> {
286 model: Option<M>,
287 normalizer: Option<N>,
288 pre_tokenizer: Option<PT>,
289 post_processor: Option<PP>,
290 decoder: Option<D>,
291
292 added_vocabulary: AddedVocabulary,
293
294 truncation: Option<TruncationParams>,
295 padding: Option<PaddingParams>,
296}
297
298impl<M, N, PT, PP, D> Default for TokenizerBuilder<M, N, PT, PP, D>
299where
300 M: Model,
301 N: Normalizer,
302 PT: PreTokenizer,
303 PP: PostProcessor,
304 D: Decoder,
305{
306 fn default() -> Self {
307 Self::new()
308 }
309}
310
311impl<M, N, PT, PP, D> TokenizerBuilder<M, N, PT, PP, D>
312where
313 M: Model,
314 N: Normalizer,
315 PT: PreTokenizer,
316 PP: PostProcessor,
317 D: Decoder,
318{
319 pub fn new() -> Self {
321 Self {
322 model: None,
323 normalizer: None,
324 pre_tokenizer: None,
325 post_processor: None,
326 decoder: None,
327 added_vocabulary: AddedVocabulary::new(),
328 truncation: None,
329 padding: None,
330 }
331 }
332
333 pub fn build(self) -> Result<TokenizerImpl<M, N, PT, PP, D>> {
337 let model = self
338 .model
339 .ok_or_else(|| Box::new(BuilderError("Model missing.".into())))?;
340 Ok(TokenizerImpl {
341 normalizer: self.normalizer,
342 pre_tokenizer: self.pre_tokenizer,
343 model,
344
345 post_processor: self.post_processor,
346 decoder: self.decoder,
347 added_vocabulary: self.added_vocabulary,
348 truncation: self.truncation,
349 padding: self.padding,
350 })
351 }
352
353 #[must_use]
355 pub fn with_model(mut self, model: M) -> Self {
356 self.model = Some(model);
357 self
358 }
359
360 #[must_use]
362 pub fn with_normalizer(mut self, normalizer: Option<N>) -> Self {
363 self.normalizer = normalizer;
364 self
365 }
366
367 #[must_use]
369 pub fn with_pre_tokenizer(mut self, pretokenizer: Option<PT>) -> Self {
370 self.pre_tokenizer = pretokenizer;
371 self
372 }
373
374 #[must_use]
376 pub fn with_post_processor(mut self, post_processor: Option<PP>) -> Self {
377 self.post_processor = post_processor;
378 self
379 }
380
381 #[must_use]
383 pub fn with_decoder(mut self, decoder: Option<D>) -> Self {
384 self.decoder = decoder;
385 self
386 }
387
388 pub fn with_added_vocabulary(mut self, added_vocabulary: AddedVocabulary) -> Self {
390 self.added_vocabulary = added_vocabulary;
391 self
392 }
393
394 #[must_use]
396 pub fn with_truncation(mut self, trunc: Option<TruncationParams>) -> Self {
397 self.truncation = trunc;
398 self
399 }
400
401 #[must_use]
403 pub fn with_padding(mut self, padding: Option<PaddingParams>) -> Self {
404 self.padding = padding;
405 self
406 }
407}
408
409#[derive(Serialize, Deserialize, Debug, Clone)]
410pub struct Tokenizer(
411 TokenizerImpl<
412 ModelWrapper,
413 NormalizerWrapper,
414 PreTokenizerWrapper,
415 PostProcessorWrapper,
416 DecoderWrapper,
417 >,
418);
419
420impl Tokenizer {
421 pub fn new(model: impl Into<ModelWrapper>) -> Self {
423 Self(TokenizerImpl::new(model.into()))
424 }
425
426 pub fn into_inner(
428 self,
429 ) -> TokenizerImpl<
430 ModelWrapper,
431 NormalizerWrapper,
432 PreTokenizerWrapper,
433 PostProcessorWrapper,
434 DecoderWrapper,
435 > {
436 self.0
437 }
438 pub fn from_file<P: AsRef<Path>>(file: P) -> Result<Self> {
439 let content = read_to_string(file)?;
440 let tokenizer = serde_json::from_str(&content)?;
441 Ok(tokenizer)
442 }
443 pub fn from_bytes<P: AsRef<[u8]>>(bytes: P) -> Result<Self> {
444 let tokenizer = serde_json::from_slice(bytes.as_ref())?;
445 Ok(tokenizer)
446 }
447 #[cfg(feature = "http")]
448 pub fn from_pretrained<S: AsRef<str>>(
449 identifier: S,
450 params: Option<crate::utils::from_pretrained::FromPretrainedParameters>,
451 ) -> Result<Self> {
452 let tokenizer_file = crate::utils::from_pretrained::from_pretrained(identifier, params)?;
453 Tokenizer::from_file(tokenizer_file)
454 }
455}
456
457impl std::str::FromStr for Tokenizer {
458 type Err = Box<dyn std::error::Error + Send + Sync>;
459
460 fn from_str(s: &str) -> Result<Self> {
461 Ok(serde_json::from_str(s)?)
462 }
463}
464
465impl<M, N, PT, PP, D> From<TokenizerImpl<M, N, PT, PP, D>> for Tokenizer
466where
467 M: Into<ModelWrapper>,
468 N: Into<NormalizerWrapper>,
469 PT: Into<PreTokenizerWrapper>,
470 PP: Into<PostProcessorWrapper>,
471 D: Into<DecoderWrapper>,
472{
473 fn from(t: TokenizerImpl<M, N, PT, PP, D>) -> Self {
474 Self(TokenizerImpl {
475 model: t.model.into(),
476 normalizer: t.normalizer.map(Into::into),
477 pre_tokenizer: t.pre_tokenizer.map(Into::into),
478 post_processor: t.post_processor.map(Into::into),
479 decoder: t.decoder.map(Into::into),
480 added_vocabulary: t.added_vocabulary,
481 padding: t.padding,
482 truncation: t.truncation,
483 })
484 }
485}
486
487impl Deref for Tokenizer {
488 type Target = TokenizerImpl<
489 ModelWrapper,
490 NormalizerWrapper,
491 PreTokenizerWrapper,
492 PostProcessorWrapper,
493 DecoderWrapper,
494 >;
495
496 fn deref(&self) -> &Self::Target {
497 &self.0
498 }
499}
500
501impl DerefMut for Tokenizer {
502 fn deref_mut(&mut self) -> &mut Self::Target {
503 &mut self.0
504 }
505}
506
507#[derive(thiserror::Error, Debug)]
508#[error("{0}")]
509pub struct TruncationParamError(String);
510
511#[derive(Clone, Debug)]
513pub struct TokenizerImpl<M, N, PT, PP, D> {
514 normalizer: Option<N>,
516 pre_tokenizer: Option<PT>,
517 model: M,
518 post_processor: Option<PP>,
519 decoder: Option<D>,
520
521 added_vocabulary: AddedVocabulary,
523
524 truncation: Option<TruncationParams>,
526 padding: Option<PaddingParams>,
527}
528
529impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
530where
531 M: Model,
532 N: Normalizer,
533 PT: PreTokenizer,
534 PP: PostProcessor,
535 D: Decoder,
536{
537 pub fn new(model: M) -> Self {
539 Self {
540 normalizer: None,
541 pre_tokenizer: None,
542 model,
543 post_processor: None,
544 decoder: None,
545
546 added_vocabulary: AddedVocabulary::new(),
547
548 truncation: None,
549 padding: None,
550 }
551 }
552
553 pub fn with_normalizer(&mut self, normalizer: Option<impl Into<N>>) -> &mut Self {
555 self.normalizer = normalizer.map(|norm| norm.into());
556 self
557 }
558 pub fn get_normalizer(&self) -> Option<&N> {
560 self.normalizer.as_ref()
561 }
562
563 pub fn with_pre_tokenizer(&mut self, pre_tokenizer: Option<impl Into<PT>>) -> &mut Self {
565 self.pre_tokenizer = pre_tokenizer.map(|tok| tok.into());
566 self
567 }
568
569 pub fn get_pre_tokenizer(&self) -> Option<&PT> {
571 self.pre_tokenizer.as_ref()
572 }
573
574 pub fn with_post_processor(&mut self, post_processor: Option<impl Into<PP>>) -> &mut Self {
576 self.post_processor = post_processor.map(|post_proc| post_proc.into());
577 self
578 }
579
580 pub fn get_post_processor(&self) -> Option<&PP> {
582 self.post_processor.as_ref()
583 }
584
585 pub fn with_decoder(&mut self, decoder: Option<impl Into<D>>) -> &mut Self {
587 self.decoder = decoder.map(|dec| dec.into());
588 self
589 }
590
591 pub fn get_decoder(&self) -> Option<&D> {
593 self.decoder.as_ref()
594 }
595
596 pub fn with_model(&mut self, model: impl Into<M>) -> &mut Self {
598 self.model = model.into();
599 self
600 }
601
602 pub fn get_model(&self) -> &M {
604 &self.model
605 }
606
607 pub fn with_added_vocabulary(&mut self, added_vocabulary: AddedVocabulary) -> &mut Self {
609 self.added_vocabulary = added_vocabulary;
610 self
611 }
612
613 pub fn get_added_vocabulary(&self) -> &AddedVocabulary {
615 &self.added_vocabulary
616 }
617
618 pub fn with_truncation(&mut self, trunc: Option<TruncationParams>) -> Result<&mut Self> {
622 if let Some(trunc_params) = &trunc {
623 let n_added_tokens = self.get_n_added_tokens(false);
624 let effective_max_length = trunc_params.max_length - n_added_tokens;
625 if effective_max_length < trunc_params.stride {
626 return Err(Box::new(TruncationParamError(format!(
627 "tokenizer stride set to {}, which is greater than or equal to its effective max length of {} (= {} original max length - {} added special tokens), ",
628 trunc_params.stride, effective_max_length, trunc_params.max_length, n_added_tokens
629 ))));
630 }
631 }
632 self.truncation = trunc;
633 Ok(self)
634 }
635
636 pub fn get_truncation(&self) -> Option<&TruncationParams> {
638 self.truncation.as_ref()
639 }
640
641 pub fn get_truncation_mut(&mut self) -> Option<&mut TruncationParams> {
643 self.truncation.as_mut()
644 }
645
646 pub fn with_padding(&mut self, padding: Option<PaddingParams>) -> &mut Self {
648 self.padding = padding;
649 self
650 }
651
652 pub fn get_padding(&self) -> Option<&PaddingParams> {
654 self.padding.as_ref()
655 }
656
657 pub fn get_padding_mut(&mut self) -> Option<&mut PaddingParams> {
659 self.padding.as_mut()
660 }
661
662 pub fn get_vocab(&self, with_added_tokens: bool) -> HashMap<String, u32> {
664 let mut final_vocab = self.model.get_vocab();
665
666 if with_added_tokens {
667 let added_vocab = self.added_vocabulary.get_vocab();
668 if !added_vocab.is_empty() {
669 final_vocab.reserve(added_vocab.len());
670 for (token, id) in added_vocab {
671 final_vocab.insert(token.clone(), *id);
672 }
673 }
674 }
675
676 final_vocab
677 }
678
679 pub fn get_added_tokens_decoder(&self) -> AHashMap<u32, AddedToken> {
681 self.added_vocabulary.get_added_tokens_decoder().clone()
682 }
683
684 pub fn get_vocab_size(&self, with_added_tokens: bool) -> usize {
686 if with_added_tokens {
689 self.get_vocab(true).len()
690 } else {
691 self.model.get_vocab_size()
692 }
693 }
694
695 pub fn token_to_id(&self, token: &str) -> Option<u32> {
697 self.added_vocabulary.token_to_id(token, &self.model)
698 }
699
700 pub fn id_to_token(&self, id: u32) -> Option<String> {
702 self.added_vocabulary
703 .simple_id_to_token(id)
704 .or_else(|| self.model.id_to_token(id))
705 }
706
707 pub fn set_encode_special_tokens(&mut self, value: bool) {
709 self.added_vocabulary.set_encode_special_tokens(value);
710 }
711
712 pub fn get_encode_special_tokens(&self) -> bool {
714 self.added_vocabulary.get_encode_special_tokens()
715 }
716
717 fn encode_single_sequence(
719 &self,
720 sequence: InputSequence,
721 type_id: u32,
722 offsets_type: OffsetType,
723 ) -> Result<Encoding> {
724 let encode = |is_pre_tokenized, subseq_idx, subseq| -> Result<Encoding> {
725 let normalized = self
726 .added_vocabulary
727 .extract_and_normalize(self.normalizer.as_ref(), subseq);
728 let pre_tokenized = self.do_pre_tokenize(normalized)?;
729 let subseq_encoding = self.do_tokenize(
730 pre_tokenized,
731 type_id,
732 if is_pre_tokenized {
733 Some(subseq_idx as u32)
734 } else {
735 None
736 },
737 offsets_type,
738 )?;
739
740 Ok(subseq_encoding)
741 };
742
743 match sequence {
744 InputSequence::PreTokenized(seq) => seq
745 .iter()
746 .enumerate()
747 .map(|(i, sequence)| encode(true, i, sequence))
748 .collect(),
749 InputSequence::PreTokenizedOwned(seq) => seq
750 .iter()
751 .enumerate()
752 .map(|(i, sequence)| encode(true, i, sequence))
753 .collect(),
754 InputSequence::PreTokenizedCow(seq) => seq
755 .iter()
756 .enumerate()
757 .map(|(i, sequence)| encode(true, i, sequence))
758 .collect(),
759 InputSequence::Raw(seq) => encode(false, 0, seq.as_ref()),
760 }
761 }
762
763 pub fn encode_fast<'s, E>(&self, input: E, add_special_tokens: bool) -> Result<Encoding>
786 where
787 E: Into<EncodeInput<'s>>,
788 {
789 let (sequence, pair) = match input.into() {
791 EncodeInput::Single(s1) => (s1, None),
792 EncodeInput::Dual(s1, s2) => (s1, Some(s2)),
793 };
794
795 let encoding = self.encode_single_sequence(sequence, 0, OffsetType::None)?;
797 let pair_encoding = pair
798 .map(|sequence| self.encode_single_sequence(sequence, 1, OffsetType::None))
799 .transpose()?;
800
801 self.post_process(encoding, pair_encoding, add_special_tokens)
803 }
804
805 pub fn encode<'s, E>(&self, input: E, add_special_tokens: bool) -> Result<Encoding>
828 where
829 E: Into<EncodeInput<'s>>,
830 {
831 let (sequence, pair) = match input.into() {
833 EncodeInput::Single(s1) => (s1, None),
834 EncodeInput::Dual(s1, s2) => (s1, Some(s2)),
835 };
836
837 let encoding = self.encode_single_sequence(sequence, 0, OffsetType::Byte)?;
839 let pair_encoding = pair
840 .map(|sequence| self.encode_single_sequence(sequence, 1, OffsetType::Byte))
841 .transpose()?;
842
843 self.post_process(encoding, pair_encoding, add_special_tokens)
845 }
846
847 pub fn encode_char_offsets<'s, E>(&self, input: E, add_special_tokens: bool) -> Result<Encoding>
871 where
872 E: Into<EncodeInput<'s>>,
873 {
874 let (sequence, pair) = match input.into() {
876 EncodeInput::Single(s1) => (s1, None),
877 EncodeInput::Dual(s1, s2) => (s1, Some(s2)),
878 };
879
880 let encoding = self.encode_single_sequence(sequence, 0, OffsetType::Char)?;
882 let pair_encoding = pair
883 .map(|sequence| self.encode_single_sequence(sequence, 1, OffsetType::Char))
884 .transpose()?;
885
886 self.post_process(encoding, pair_encoding, add_special_tokens)
888 }
889
890 pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result<String> {
892 let tokens = ids
893 .iter()
894 .filter_map(|id| {
895 self.added_vocabulary
896 .simple_id_to_token(*id)
897 .or_else(|| self.model.id_to_token(*id))
898 .filter(|token| {
899 !skip_special_tokens || !self.added_vocabulary.is_special_token(token)
900 })
901 })
902 .collect::<Vec<_>>();
903
904 if let Some(decoder) = &self.decoder {
905 decoder.decode(tokens)
906 } else {
907 Ok(tokens.join(" "))
908 }
909 }
910
911 pub fn decode_stream(&self, skip_special_tokens: bool) -> DecodeStream<'_, M, N, PT, PP, D> {
914 DecodeStream::new(self, skip_special_tokens)
915 }
916}
917
918pub struct DecodeStream<'tok, M, N, PT, PP, D> {
1018 tokenizer: &'tok TokenizerImpl<M, N, PT, PP, D>,
1020 skip_special_tokens: bool,
1022 ids: Vec<u32>,
1034 prefix: String,
1037 prefix_index: usize,
1040}
1041
1042#[derive(thiserror::Error, Debug)]
1043pub enum DecodeStreamError {
1044 #[error("Invalid prefix encountered")]
1045 InvalidPrefix,
1046}
1047
1048impl<'tok, M, N, PT, PP, D> DecodeStream<'tok, M, N, PT, PP, D>
1049where
1050 M: Model,
1051 N: Normalizer,
1052 PT: PreTokenizer,
1053 PP: PostProcessor,
1054 D: Decoder,
1055{
1056 fn new(tokenizer: &'tok TokenizerImpl<M, N, PT, PP, D>, skip_special_tokens: bool) -> Self {
1057 Self {
1058 tokenizer,
1059 ids: vec![],
1060 skip_special_tokens,
1061 prefix: "".to_string(),
1062 prefix_index: 0,
1063 }
1064 }
1065
1066 pub fn step(&mut self, id: u32) -> Result<Option<String>> {
1068 step_decode_stream(
1069 self.tokenizer,
1070 id,
1071 self.skip_special_tokens,
1072 &mut self.ids,
1073 &mut self.prefix,
1074 &mut self.prefix_index,
1075 )
1076 }
1077}
1078
1079pub fn step_decode_stream<M, N, PT, PP, D>(
1081 tokenizer: &TokenizerImpl<M, N, PT, PP, D>,
1082 id: u32,
1083 skip_special_tokens: bool,
1084 ids: &mut Vec<u32>,
1085 prefix: &mut String,
1086 prefix_index: &mut usize,
1087) -> Result<Option<String>>
1088where
1089 M: Model,
1090 N: Normalizer,
1091 PT: PreTokenizer,
1092 PP: PostProcessor,
1093 D: Decoder,
1094{
1095 ids.push(id);
1096 let string = tokenizer.decode(ids.as_slice(), skip_special_tokens)?;
1097 if string.len() > prefix.len() && !string.ends_with('�') {
1098 if !(string.starts_with(&*prefix)) {
1099 return Err(Box::new(DecodeStreamError::InvalidPrefix));
1100 }
1101 let new_text = &string[prefix.len()..].to_string();
1102 let new_prefix_index = ids.len() - *prefix_index;
1103 *ids = ids.drain(*prefix_index..).collect();
1104 *prefix = tokenizer.decode(ids, skip_special_tokens)?;
1105 *prefix_index = new_prefix_index;
1106 Ok(Some(new_text.to_string()))
1107 } else {
1108 Ok(None)
1109 }
1110}
1111
1112impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
1113where
1114 M: Model,
1115{
1116 fn do_tokenize<P: Into<PreTokenizedString>>(
1119 &self,
1120 pretokenized: P,
1121 type_id: u32,
1122 word_idx: Option<u32>,
1123 offsets_type: OffsetType,
1124 ) -> Result<Encoding> {
1125 let mut pretokenized: PreTokenizedString = pretokenized.into();
1126 pretokenized.tokenize(|normalized| self.model.tokenize(normalized.get()))?;
1127 pretokenized.into_encoding(word_idx, type_id, offsets_type)
1128 }
1129}
1130
1131#[allow(dead_code)]
1132impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
1133where
1134 N: Normalizer,
1135{
1136 fn do_normalize<V: Into<NormalizedString>>(&self, normalized: V) -> Result<NormalizedString> {
1138 let mut normalized: NormalizedString = normalized.into();
1139
1140 if let Some(ref normalizer) = self.normalizer {
1141 normalizer.normalize(&mut normalized)?;
1142 }
1143
1144 Ok(normalized)
1145 }
1146}
1147
1148impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
1149where
1150 N: Normalizer,
1151 M: Model,
1152{
1153 pub fn add_special_tokens(&mut self, tokens: &[AddedToken]) -> usize {
1156 self.added_vocabulary
1157 .add_special_tokens(tokens, &self.model, self.normalizer.as_ref())
1158 }
1159
1160 pub fn add_tokens(&mut self, tokens: &[AddedToken]) -> usize {
1162 self.added_vocabulary
1163 .add_tokens(tokens, &self.model, self.normalizer.as_ref())
1164 }
1165}
1166
1167impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
1168where
1169 PT: PreTokenizer,
1170{
1171 fn do_pre_tokenize<P: Into<PreTokenizedString>>(
1173 &self,
1174 pretokenized: P,
1175 ) -> Result<PreTokenizedString> {
1176 let mut pretokenized: PreTokenizedString = pretokenized.into();
1177 if let Some(ref pretok) = self.pre_tokenizer {
1178 pretok.pre_tokenize(&mut pretokenized)?;
1179 }
1180
1181 Ok(pretokenized)
1182 }
1183}
1184
1185impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
1186where
1187 PP: PostProcessor,
1188{
1189 pub fn post_process(
1191 &self,
1192 encoding: Encoding,
1193 pair_encoding: Option<Encoding>,
1194 add_special_tokens: bool,
1195 ) -> Result<Encoding> {
1196 let (encoding, pair_encoding) = {
1198 if let Some(trunc) = &self.truncation {
1199 let n_added_tokens = self.get_n_added_tokens(pair_encoding.is_some());
1200
1201 if add_special_tokens && n_added_tokens > 0 {
1202 let params = TruncationParams {
1203 max_length: trunc.max_length - n_added_tokens,
1204 ..*trunc
1205 };
1206 truncate_encodings(encoding, pair_encoding, ¶ms)?
1207 } else {
1208 truncate_encodings(encoding, pair_encoding, trunc)?
1209 }
1210 } else {
1211 (encoding, pair_encoding)
1212 }
1213 };
1214
1215 let final_encoding = if let Some(processor) = &self.post_processor {
1217 processor.process(encoding, pair_encoding, add_special_tokens)?
1218 } else {
1219 let encodings = if let Some(pair_encoding) = pair_encoding {
1220 vec![encoding, pair_encoding]
1221 } else {
1222 vec![encoding]
1223 };
1224 let mut encodings =
1225 <dyn PostProcessor>::default_process(encodings, add_special_tokens)?;
1226 if encodings.len() != 1 {
1227 panic!("We haven't reduced the encodings like we should have");
1228 }
1229 encodings.pop().unwrap()
1230 };
1231
1232 let [final_encoding] = if let Some(params) = &self.padding {
1234 let mut arr = [final_encoding];
1235 pad_encodings(&mut arr, params)?;
1236 arr
1237 } else {
1238 [final_encoding]
1239 };
1240
1241 Ok(final_encoding)
1242 }
1243
1244 fn get_n_added_tokens(&self, is_pair: bool) -> usize {
1245 if let Some(processor) = &self.post_processor {
1246 processor.added_tokens(is_pair)
1247 } else {
1248 0
1249 }
1250 }
1251}
1252
1253impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
1254where
1255 M: Model + Send + Sync,
1256 N: Normalizer + Send + Sync,
1257 PT: PreTokenizer + Send + Sync,
1258 PP: PostProcessor + Send + Sync,
1259 D: Decoder + Send + Sync,
1260{
1261 pub fn encode_batch<'s, E>(
1263 &self,
1264 inputs: Vec<E>,
1265 add_special_tokens: bool,
1266 ) -> Result<Vec<Encoding>>
1267 where
1268 E: Into<EncodeInput<'s>> + Send,
1269 {
1270 let mut encodings = inputs
1271 .into_maybe_par_iter()
1272 .map(|input| self.encode(input, add_special_tokens))
1273 .collect::<Result<Vec<Encoding>>>()?;
1274
1275 if let Some(params) = &self.padding {
1276 pad_encodings(&mut encodings, params)?;
1278 }
1279
1280 Ok(encodings)
1281 }
1282
1283 pub fn encode_batch_char_offsets<'s, E>(
1286 &self,
1287 inputs: Vec<E>,
1288 add_special_tokens: bool,
1289 ) -> Result<Vec<Encoding>>
1290 where
1291 E: Into<EncodeInput<'s>> + Send,
1292 {
1293 let mut encodings = inputs
1294 .into_maybe_par_iter()
1295 .map(|input| self.encode_char_offsets(input, add_special_tokens))
1296 .collect::<Result<Vec<Encoding>>>()?;
1297
1298 if let Some(params) = &self.padding {
1299 pad_encodings(&mut encodings, params)?;
1301 }
1302
1303 Ok(encodings)
1304 }
1305
1306 pub fn encode_batch_fast<'s, E>(
1308 &self,
1309 inputs: Vec<E>,
1310 add_special_tokens: bool,
1311 ) -> Result<Vec<Encoding>>
1312 where
1313 E: Into<EncodeInput<'s>> + Send,
1314 {
1315 let mut encodings = inputs
1316 .into_maybe_par_iter()
1317 .map(|input| self.encode_fast(input, add_special_tokens))
1318 .collect::<Result<Vec<Encoding>>>()?;
1319
1320 if let Some(params) = &self.padding {
1321 pad_encodings(&mut encodings, params)?;
1323 }
1324
1325 Ok(encodings)
1326 }
1327
1328 pub fn decode_batch(
1330 &self,
1331 sentences: &[&[u32]],
1332 skip_special_tokens: bool,
1333 ) -> Result<Vec<String>>
1334 where
1335 M: Send + Sync,
1336 {
1337 sentences
1338 .into_maybe_par_iter()
1339 .map(|sentence| self.decode(sentence, skip_special_tokens))
1340 .collect()
1341 }
1342
1343 pub fn train_from_files<T>(&mut self, trainer: &mut T, files: Vec<String>) -> Result<&mut Self>
1345 where
1346 T: Trainer<Model = M> + Sync,
1347 {
1348 let mut len = 0;
1349 for file in files.iter() {
1350 len += File::open(file)
1351 .and_then(|f| f.metadata())
1352 .map(|m| m.len())?;
1353 }
1354
1355 let max_read = 1_000_000;
1356
1357 ResultShunt::process(
1358 files.into_iter().flat_map(|filename| {
1359 match File::open(filename) {
1360 Ok(file) => {
1361 let file = BufReader::with_capacity(max_read, file);
1362 itertools::Either::Left(file.lines_with_ending())
1366 }
1367 Err(e) => itertools::Either::Right(std::iter::once(Err(e))),
1368 }
1369 }),
1370 |sequences| -> Result<()> {
1371 let progress = if trainer.should_show_progress() {
1372 let progress = ProgressBar::new(len);
1373 progress.set_style(
1374 ProgressStyle::default_bar()
1375 .template("[{elapsed_precise}] {msg:<30!} {wide_bar} {percent:>18!}%")
1376 .expect("Invalid progress template"),
1377 );
1378 progress
1379 .set_message(format!("Pre-processing files ({:.2} Mo)", len / 1_000_000));
1380 Some(progress)
1381 } else {
1382 None
1383 };
1384
1385 trainer.feed(
1386 sequences.inspect(|s| {
1387 if let Some(progress) = &progress {
1388 progress.inc(s.len() as u64)
1389 }
1390 }),
1391 |seq| {
1392 let normalized = self
1393 .added_vocabulary
1394 .extract_and_normalize(self.normalizer.as_ref(), seq.as_ref());
1395 let pre_tokenized = self.do_pre_tokenize(normalized)?;
1396 Ok(pre_tokenized
1397 .get_splits(OffsetReferential::Original, OffsetType::Byte)
1398 .into_iter()
1399 .map(|(s, _, _)| s.to_owned())
1400 .collect())
1401 },
1402 )?;
1403
1404 if let Some(pbar) = progress {
1405 pbar.finish();
1406 }
1407 let special_tokens = trainer.train(&mut self.model)?;
1408 self.add_special_tokens(&special_tokens);
1409
1410 Ok(())
1411 },
1412 )??;
1413 Ok(self)
1414 }
1415
1416 pub fn train<T, I, S>(&mut self, trainer: &mut T, sequences: I) -> Result<&mut Self>
1418 where
1419 T: Trainer<Model = M> + Sync,
1420 I: Iterator<Item = S> + Send,
1421 S: AsRef<str> + Send,
1422 {
1423 let (lower, upper) = sequences.size_hint();
1424 let len = upper.unwrap_or(lower) as u64;
1425 let progress = if trainer.should_show_progress() {
1426 let progress = ProgressBar::new(len);
1427 progress.set_style(
1428 ProgressStyle::default_bar()
1429 .template("[{elapsed_precise}] {msg:<30!} {wide_bar} {pos:<9!}/{len:>9!}")
1430 .expect("Invalid progress template"),
1431 );
1432 progress.set_message("Pre-processing sequences");
1433 Some(progress)
1434 } else {
1435 None
1436 };
1437
1438 trainer.feed(
1439 sequences.inspect(|_s| {
1440 if let Some(progress) = &progress {
1441 progress.inc(1)
1442 }
1443 }),
1444 |seq| {
1445 let normalized = self
1446 .added_vocabulary
1447 .extract_and_normalize(self.normalizer.as_ref(), seq.as_ref());
1448 let pre_tokenized = self.do_pre_tokenize(normalized)?;
1449 Ok(pre_tokenized
1450 .get_splits(OffsetReferential::Original, OffsetType::Byte)
1451 .into_iter()
1452 .map(|(s, _, _)| s.to_owned())
1453 .collect())
1454 },
1455 )?;
1456 if let Some(pbar) = progress {
1457 pbar.finish();
1458 }
1459
1460 let special_tokens = trainer.train(&mut self.model)?;
1461 self.add_special_tokens(&special_tokens);
1462
1463 Ok(self)
1464 }
1465}
1466
1467impl<M, N, PT, PP, D> std::str::FromStr for TokenizerImpl<M, N, PT, PP, D>
1468where
1469 M: for<'de> Deserialize<'de> + Model,
1470 N: for<'de> Deserialize<'de> + Normalizer,
1471 PT: for<'de> Deserialize<'de> + PreTokenizer,
1472 PP: for<'de> Deserialize<'de> + PostProcessor,
1473 D: for<'de> Deserialize<'de> + Decoder,
1474{
1475 type Err = Error;
1476
1477 fn from_str(s: &str) -> Result<Self> {
1478 Ok(serde_json::from_str(s)?)
1479 }
1480}
1481
1482impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
1483where
1484 M: DeserializeOwned + Model,
1485 N: DeserializeOwned + Normalizer,
1486 PT: DeserializeOwned + PreTokenizer,
1487 PP: DeserializeOwned + PostProcessor,
1488 D: DeserializeOwned + Decoder,
1489{
1490 pub fn from_file<P: AsRef<Path>>(file: P) -> Result<Self> {
1492 let content = read_to_string(file)?;
1493 let tokenizer = serde_json::from_str(&content)?;
1494 Ok(tokenizer)
1495 }
1496}
1497
1498impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
1499where
1500 M: DeserializeOwned + Model,
1501 N: DeserializeOwned + Normalizer,
1502 PT: DeserializeOwned + PreTokenizer,
1503 PP: DeserializeOwned + PostProcessor,
1504 D: DeserializeOwned + Decoder,
1505{
1506 pub fn from_bytes<P: AsRef<[u8]>>(bytes: P) -> Result<Self> {
1508 let tokenizer = serde_json::from_slice(bytes.as_ref())?;
1509 Ok(tokenizer)
1510 }
1511}
1512
1513impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
1514where
1515 M: DeserializeOwned + Model,
1516 N: DeserializeOwned + Normalizer,
1517 PT: DeserializeOwned + PreTokenizer,
1518 PP: DeserializeOwned + PostProcessor,
1519 D: DeserializeOwned + Decoder,
1520{
1521 #[deprecated(
1522 since = "0.14.0",
1523 note = "Users should download the file separately using https://github.com/huggingface/hf-hub instead, which splits concerns of accessing the web, and should use the new cache layout"
1524 )]
1525 #[cfg(feature = "http")]
1526 pub fn from_pretrained<S: AsRef<str>>(
1529 identifier: S,
1530 params: Option<crate::utils::from_pretrained::FromPretrainedParameters>,
1531 ) -> Result<Self> {
1532 let tokenizer_file = crate::utils::from_pretrained::from_pretrained(identifier, params)?;
1533 TokenizerImpl::from_file(tokenizer_file)
1534 }
1535}
1536
1537impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>
1538where
1539 M: Serialize,
1540 N: Serialize,
1541 PT: Serialize,
1542 PP: Serialize,
1543 D: Serialize,
1544{
1545 pub fn to_string(&self, pretty: bool) -> Result<String> {
1547 Ok(if pretty {
1548 serde_json::to_string_pretty(self)?
1549 } else {
1550 serde_json::to_string(self)?
1551 })
1552 }
1553
1554 pub fn save<P: AsRef<Path>>(&self, path: P, pretty: bool) -> Result<()> {
1556 let serialized = self.to_string(pretty)?;
1557
1558 let mut file = File::create(path)?;
1559 file.write_all(serialized.as_bytes())?;
1560
1561 Ok(())
1562 }
1563}