tokenizers/tokenizer/
mod.rs

1//! Represents a tokenization pipeline.
2//!
3//! A [`Tokenizer`](struct.Tokenizer.html) is composed of some of the following parts.
4//!   - [`Normalizer`](trait.Normalizer.html): Takes care of the text normalization (like unicode normalization).
5//!   - [`PreTokenizer`](trait.PreTokenizer.html): Takes care of the pre tokenization (ie. How to split tokens and pre-process
6//!     them.
7//!   - [`Model`](trait.Model.html): A model encapsulates the tokenization algorithm (like BPE, Word base, character
8//!     based, ...).
9//!   - [`PostProcessor`](trait.PostProcessor.html): Takes care of the processing after tokenization (like truncating, padding,
10//!     ...).
11
12use 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
34// Re-export wrappers
35pub 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;
40// And some other types
41pub 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
55/// Takes care of pre-processing strings.
56pub trait Normalizer {
57    fn normalize(&self, normalized: &mut NormalizedString) -> Result<()>;
58}
59
60/// The `PreTokenizer` is in charge of doing the pre-segmentation step. It splits the given string
61/// in multiple substrings, keeping track of the offsets of said substrings from the
62/// `NormalizedString`. In some occasions, the `PreTokenizer` might need to modify the given
63/// `NormalizedString` to ensure we can entirely keep track of the offsets and the mapping with
64/// the original string.
65pub trait PreTokenizer {
66    fn pre_tokenize(&self, pretokenized: &mut PreTokenizedString) -> Result<()>;
67}
68
69/// Represents a model used during Tokenization (like BPE or Word or Unigram).
70pub trait Model {
71    type Trainer: Trainer + Sync;
72    /// Tokenize the given sequence into multiple underlying `Token`. The `offsets` on the `Token`
73    /// are expected to be relative to the given sequence.
74    fn tokenize(&self, sequence: &str) -> Result<Vec<Token>>;
75    /// Find the ID associated to a string token
76    fn token_to_id(&self, token: &str) -> Option<u32>;
77    /// Find the string token associated to an ID
78    fn id_to_token(&self, id: u32) -> Option<String>;
79    /// Retrieve the entire vocabulary mapping (token -> ID)
80    fn get_vocab(&self) -> HashMap<String, u32>;
81    /// Retrieve the size of the vocabulary
82    fn get_vocab_size(&self) -> usize;
83    /// Save the current `Model` in the given folder, using the given `prefix` for the various
84    /// files that need to be saved.
85    fn save(&self, folder: &Path, prefix: Option<&str>) -> Result<Vec<PathBuf>>;
86    /// Get an instance of a Trainer capable of training this Model
87    fn get_trainer(&self) -> <Self as Model>::Trainer;
88}
89
90/// A `PostProcessor` has the responsibility to post process an encoded output of the `Tokenizer`.
91/// It adds any special tokens that a language model would require.
92pub trait PostProcessor {
93    /// Returns the number of tokens that will be added during the processing step
94    fn added_tokens(&self, is_pair: bool) -> usize;
95    /// Process both encodings and returns a new merged one
96    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    /// Process any amount of encodings and returns a series of encoding (might merge them)
121    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
152/// A `Decoder` changes the raw tokens into its more readable form.
153pub 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
161/// A `Trainer` has the responsibility to train a model. We feed it with lines/sentences
162/// and then it can train the given `Model`.
163pub trait Trainer {
164    type Model: Model + Sized;
165    /// Whether we should show progress during the training.
166    fn should_show_progress(&self) -> bool;
167    /// The actual training method. This will return a new trained Model as well as a list
168    /// of `special_tokens` to be added directly to the tokenizer along with the model.
169    fn train(&self, model: &mut Self::Model) -> Result<Vec<AddedToken>>;
170    /// Process an iterator of sequences, calling `process` for each of them in order to
171    /// pre-process the said sequence as relevant.
172    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
282/// Builder for Tokenizer structs.
283///
284/// `build()` fails if the `model` is missing.
285pub 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    /// Get an empty TokenizerBuilder.
320    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    /// Convert the TokenizerBuilder to a Tokenizer.
334    ///
335    /// Conversion fails if the `model` is missing.
336    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    /// Set the model.
354    #[must_use]
355    pub fn with_model(mut self, model: M) -> Self {
356        self.model = Some(model);
357        self
358    }
359
360    /// Set the normalizer.
361    #[must_use]
362    pub fn with_normalizer(mut self, normalizer: Option<N>) -> Self {
363        self.normalizer = normalizer;
364        self
365    }
366
367    /// Set the pre-tokenizer.
368    #[must_use]
369    pub fn with_pre_tokenizer(mut self, pretokenizer: Option<PT>) -> Self {
370        self.pre_tokenizer = pretokenizer;
371        self
372    }
373
374    /// Set the post-processor.
375    #[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    /// Set the decoder.
382    #[must_use]
383    pub fn with_decoder(mut self, decoder: Option<D>) -> Self {
384        self.decoder = decoder;
385        self
386    }
387
388    /// Set the added vocabulary.
389    pub fn with_added_vocabulary(mut self, added_vocabulary: AddedVocabulary) -> Self {
390        self.added_vocabulary = added_vocabulary;
391        self
392    }
393
394    /// Set the truncation parameters.
395    #[must_use]
396    pub fn with_truncation(mut self, trunc: Option<TruncationParams>) -> Self {
397        self.truncation = trunc;
398        self
399    }
400
401    /// Set the padding parameters.
402    #[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    /// Construct a new Tokenizer based on the model.
422    pub fn new(model: impl Into<ModelWrapper>) -> Self {
423        Self(TokenizerImpl::new(model.into()))
424    }
425
426    /// Unwrap the TokenizerImpl.
427    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/// A `Tokenizer` is capable of encoding/decoding any text.
512#[derive(Clone, Debug)]
513pub struct TokenizerImpl<M, N, PT, PP, D> {
514    // Tokenizer parts
515    normalizer: Option<N>,
516    pre_tokenizer: Option<PT>,
517    model: M,
518    post_processor: Option<PP>,
519    decoder: Option<D>,
520
521    // Added Vocabulary capabilities
522    added_vocabulary: AddedVocabulary,
523
524    // General processing parameters
525    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    /// Instantiate a new Tokenizer, with the given Model
538    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    /// Set the normalizer
554    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    /// Get the normalizer
559    pub fn get_normalizer(&self) -> Option<&N> {
560        self.normalizer.as_ref()
561    }
562
563    /// Set the pre tokenizer
564    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    /// Get the pre tokenizer
570    pub fn get_pre_tokenizer(&self) -> Option<&PT> {
571        self.pre_tokenizer.as_ref()
572    }
573
574    /// Set the post processor
575    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    /// Get the post processor
581    pub fn get_post_processor(&self) -> Option<&PP> {
582        self.post_processor.as_ref()
583    }
584
585    /// Set the decoder
586    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    /// Get the decoder
592    pub fn get_decoder(&self) -> Option<&D> {
593        self.decoder.as_ref()
594    }
595
596    /// Set the model
597    pub fn with_model(&mut self, model: impl Into<M>) -> &mut Self {
598        self.model = model.into();
599        self
600    }
601
602    /// Get the model
603    pub fn get_model(&self) -> &M {
604        &self.model
605    }
606
607    /// Set the added vocabulary.
608    pub fn with_added_vocabulary(&mut self, added_vocabulary: AddedVocabulary) -> &mut Self {
609        self.added_vocabulary = added_vocabulary;
610        self
611    }
612
613    /// Get the added vocabulary
614    pub fn get_added_vocabulary(&self) -> &AddedVocabulary {
615        &self.added_vocabulary
616    }
617
618    /// Set the truncation parameters
619    ///
620    /// Fails if `stride` is too high relative to `max_length` and `post_processor.added_tokens()`
621    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    /// Get the currently set truncation parameters
637    pub fn get_truncation(&self) -> Option<&TruncationParams> {
638        self.truncation.as_ref()
639    }
640
641    /// Get a mutable reference to the currently set truncation parameters
642    pub fn get_truncation_mut(&mut self) -> Option<&mut TruncationParams> {
643        self.truncation.as_mut()
644    }
645
646    /// Set the padding parameters
647    pub fn with_padding(&mut self, padding: Option<PaddingParams>) -> &mut Self {
648        self.padding = padding;
649        self
650    }
651
652    /// Get the currently set padding parameters
653    pub fn get_padding(&self) -> Option<&PaddingParams> {
654        self.padding.as_ref()
655    }
656
657    /// Get a mutable reference to the currently set padding parameters
658    pub fn get_padding_mut(&mut self) -> Option<&mut PaddingParams> {
659        self.padding.as_mut()
660    }
661
662    // Get the vocabulary as a plain HashMap for bindings compatibility
663    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    /// Get the added tokens decoder
680    pub fn get_added_tokens_decoder(&self) -> AHashMap<u32, AddedToken> {
681        self.added_vocabulary.get_added_tokens_decoder().clone()
682    }
683
684    /// Get the size of the vocabulary
685    pub fn get_vocab_size(&self, with_added_tokens: bool) -> usize {
686        // TODO ArthurZ THIS IS WRONG! We need to measure the length of the `set` because
687        // now some tokens can be both in the added_tokens_encoder and in the vocab
688        if with_added_tokens {
689            self.get_vocab(true).len()
690        } else {
691            self.model.get_vocab_size()
692        }
693    }
694
695    /// Converts a token in the corresponding id.
696    pub fn token_to_id(&self, token: &str) -> Option<u32> {
697        self.added_vocabulary.token_to_id(token, &self.model)
698    }
699
700    /// Converts an id to the corresponding token.
701    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    /// set the added vocab's splitting scheme
708    pub fn set_encode_special_tokens(&mut self, value: bool) {
709        self.added_vocabulary.set_encode_special_tokens(value);
710    }
711
712    /// Get added token value
713    pub fn get_encode_special_tokens(&self) -> bool {
714        self.added_vocabulary.get_encode_special_tokens()
715    }
716
717    /// Encode a single sequence
718    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    /// Encode the given input. This method accepts both single sequences, as well as pair
764    /// sequences. Also, a sequence can be a string, or already pre-tokenized input directly:
765    /// Contrarily to `encode`, it does not compute offsets
766    /// ```
767    /// # use tokenizers::Tokenizer;
768    /// # use tokenizers::models::bpe::BPE;
769    /// # let mut tokenizer = Tokenizer::new(BPE::default());
770    /// #
771    /// // Sequences:
772    /// tokenizer.encode_fast("Single sequence", false);
773    /// tokenizer.encode_fast(("Sequence A", "Sequence B"), false);
774    ///
775    /// // Pre-tokenized sequences:
776    /// tokenizer.encode_fast(&["Single", "sequence"][..], false);
777    /// tokenizer.encode_fast((
778    ///     &["Sequence", "A"][..],
779    ///     &["Sequence", "B"][..]
780    /// ), false);
781    ///
782    /// // or even both types together:
783    /// tokenizer.encode_fast(("A complete sequence", &["And", "a", "tokenized"][..]), false);
784    /// ```
785    pub fn encode_fast<'s, E>(&self, input: E, add_special_tokens: bool) -> Result<Encoding>
786    where
787        E: Into<EncodeInput<'s>>,
788    {
789        // Extract sequences from the EncodeInput
790        let (sequence, pair) = match input.into() {
791            EncodeInput::Single(s1) => (s1, None),
792            EncodeInput::Dual(s1, s2) => (s1, Some(s2)),
793        };
794
795        // Encode each sequence
796        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        // And finally post process
802        self.post_process(encoding, pair_encoding, add_special_tokens)
803    }
804
805    /// Encode the given input. This method accepts both single sequences, as well as pair
806    /// sequences. Also, a sequence can be a string, or already pre-tokenized input directly:
807    ///
808    /// ```
809    /// # use tokenizers::Tokenizer;
810    /// # use tokenizers::models::bpe::BPE;
811    /// # let mut tokenizer = Tokenizer::new(BPE::default());
812    /// #
813    /// // Sequences:
814    /// tokenizer.encode("Single sequence", false);
815    /// tokenizer.encode(("Sequence A", "Sequence B"), false);
816    ///
817    /// // Pre-tokenized sequences:
818    /// tokenizer.encode(&["Single", "sequence"][..], false);
819    /// tokenizer.encode((
820    ///     &["Sequence", "A"][..],
821    ///     &["Sequence", "B"][..]
822    /// ), false);
823    ///
824    /// // or even both types together:
825    /// tokenizer.encode(("A complete sequence", &["And", "a", "tokenized"][..]), false);
826    /// ```
827    pub fn encode<'s, E>(&self, input: E, add_special_tokens: bool) -> Result<Encoding>
828    where
829        E: Into<EncodeInput<'s>>,
830    {
831        // Extract sequences from the EncodeInput
832        let (sequence, pair) = match input.into() {
833            EncodeInput::Single(s1) => (s1, None),
834            EncodeInput::Dual(s1, s2) => (s1, Some(s2)),
835        };
836
837        // Encode each sequence
838        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        // And finally post process
844        self.post_process(encoding, pair_encoding, add_special_tokens)
845    }
846
847    /// Encode the given input, using offsets relative to chars instead of bytes.
848    /// This method accepts both single sequences, as well as pair sequences. Also,
849    /// a sequence can be a string, or already pre-tokenized input directly:
850    ///
851    /// ```
852    /// # use tokenizers::Tokenizer;
853    /// # use tokenizers::models::bpe::BPE;
854    /// # let mut tokenizer = Tokenizer::new(BPE::default());
855    /// #
856    /// // Sequences:
857    /// tokenizer.encode("Single sequence", false);
858    /// tokenizer.encode(("Sequence A", "Sequence B"), false);
859    ///
860    /// // Pre-tokenized sequences:
861    /// tokenizer.encode(&["Single", "sequence"][..], false);
862    /// tokenizer.encode((
863    ///     &["Sequence", "A"][..],
864    ///     &["Sequence", "B"][..]
865    /// ), false);
866    ///
867    /// // or even both types together:
868    /// tokenizer.encode(("A complete sequence", &["And", "a", "tokenized"][..]), false);
869    /// ```
870    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        // Extract sequences from the EncodeInput
875        let (sequence, pair) = match input.into() {
876            EncodeInput::Single(s1) => (s1, None),
877            EncodeInput::Dual(s1, s2) => (s1, Some(s2)),
878        };
879
880        // Encode each sequence
881        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        // And finally post process
887        self.post_process(encoding, pair_encoding, add_special_tokens)
888    }
889
890    /// Decode the given ids, back to a String
891    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    /// Decode the given ids, back to a String
912    /// See [`DecodeStream`]
913    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
918/// DecodeStream will keep the state necessary to produce individual chunks of
919/// strings given an input stream of token_ids.
920///
921/// This is necessary because decoding in general cannot achieve that since strings
922/// depend on surrounding ids to provide a valid string. Typically stripping extra spaces
923///
924/// Example:
925///
926/// ```
927/// # #[cfg(not(target_os = "windows"))]
928/// # {
929/// use tokenizers::Tokenizer;
930/// let tokenizer = Tokenizer::from_file("data/roberta.json").unwrap();
931///
932/// let mut decode_stream = tokenizer.decode_stream(false);
933/// assert_eq!(decode_stream.step(713).unwrap(), Some("This".to_string()));
934/// assert_eq!(decode_stream.step(16).unwrap(), Some(" is".to_string()));
935/// assert_eq!(decode_stream.step(41).unwrap(), Some(" an".to_string()));
936/// assert_eq!(
937///     decode_stream.step(1246).unwrap(),
938///     Some(" example".to_string())
939/// );
940/// # }
941/// ```
942///
943/// Returning `None` means the given id is not enough to produce a chunk.
944/// This typically happens with `byte_fallback` options where some tokens do
945/// not represent valid utf-8, and only follow-up token_ids will help produce
946/// a valid chunk.
947/// ```
948/// use tokenizers::{Tokenizer, TokenizerBuilder, models::bpe::BPE, decoders::byte_fallback::ByteFallback, pre_tokenizers::byte_level::ByteLevel, normalizers::unicode::NFC};
949/// use std::collections::HashMap;
950/// use std::iter::FromIterator;
951///
952/// let vocab = HashMap::from_iter([
953///     ("<0x20>".to_string(), 0),
954///     ("<0xC3>".to_string(), 1),
955///     ("<0xA9>".to_string(), 2),
956///     (" This".to_string(), 3),
957/// ]);
958/// let merges = vec![];
959/// let bpe = BPE::builder()
960///     .vocab_and_merges(vocab, merges)
961///     .byte_fallback(true)
962///     .build()
963///     .unwrap();
964/// let tokenizer = TokenizerBuilder::default()
965///     .with_model(bpe)
966///     .with_decoder(Some(ByteFallback::default()))
967///     .with_normalizer(Some(NFC))
968///     .with_pre_tokenizer(Some(ByteLevel::default()))
969///     .with_post_processor(Some(ByteLevel::default()))
970///     .build().unwrap();
971///
972/// let mut decode_stream = tokenizer.decode_stream(false);
973/// // Single byte_fallback is valid utf-8
974/// assert_eq!(decode_stream.step(0).unwrap(), Some(" ".to_string()));
975/// // Invalid utf-8
976/// assert_eq!(decode_stream.step(1).unwrap(), None);
977/// // Valid utf-8 again, this corresponds to both tokens: [1, 2]
978/// assert_eq!(decode_stream.step(2).unwrap(), Some("é".to_string()));
979/// ```
980///
981/// To see how [`DecodeStream`] is necessary, let's show how using raw [`TokenizerImpl::decode`] would
982/// fail.
983///
984/// ```
985/// use tokenizers::{Tokenizer, TokenizerBuilder, models::bpe::BPE, pre_tokenizers::{byte_level::ByteLevel, metaspace::Metaspace}, normalizers::unicode::NFC};
986/// use std::collections::HashMap;
987/// use std::iter::FromIterator;
988///
989/// let vocab = HashMap::from_iter([
990///     ("▁This".to_string(), 0),
991/// ]);
992/// let merges = vec![];
993/// let bpe = BPE::builder()
994///     .vocab_and_merges(vocab, merges)
995///     .byte_fallback(true)
996///     .build()
997///     .unwrap();
998/// let tokenizer = TokenizerBuilder::new()
999///     .with_model(bpe)
1000///     .with_decoder(Some(Metaspace::default()))
1001///     .with_normalizer(Some(NFC))
1002///     .with_pre_tokenizer(Some(ByteLevel::default()))
1003///     .with_post_processor(Some(ByteLevel::default()))
1004///     .build()
1005///     .unwrap();
1006///
1007/// // Strip decoder removes the extra initial space
1008/// assert_eq!(tokenizer.decode(&[0, 0], false).unwrap(), "This This");
1009/// // Decoding one token at a time would produce "ThisThis"
1010/// assert_eq!(tokenizer.decode(&[0], false).unwrap(), "This");
1011///
1012/// // Using a stream fixes it by keeping the necessary state.
1013/// let mut decode_stream = tokenizer.decode_stream(false);
1014/// assert_eq!(decode_stream.step(0).unwrap(), Some("This".to_string()));
1015/// assert_eq!(decode_stream.step(0).unwrap(), Some(" This".to_string()));
1016/// ```
1017pub struct DecodeStream<'tok, M, N, PT, PP, D> {
1018    /// A reference to the tokenizer
1019    tokenizer: &'tok TokenizerImpl<M, N, PT, PP, D>,
1020    /// Regular decode option that is kept throughout.
1021    skip_special_tokens: bool,
1022    /// A temporary buffer of the necessary token_ids needed
1023    /// to produce valid string chunks.
1024    /// This typically contains 3 parts:
1025    ///  - read
1026    ///  - prefix
1027    ///  - rest
1028    ///
1029    /// Read is the bit necessary to surround the prefix
1030    /// so decoding the whole ids produces a valid prefix.
1031    /// Prefix is the previously produced string, kept around to trim off of
1032    /// the next valid chunk
1033    ids: Vec<u32>,
1034    /// The previously returned chunk that needs to be discarded from the
1035    /// decoding of the current ids to produce the next chunk
1036    prefix: String,
1037    /// The index within the ids corresponding to the prefix so we can drain
1038    /// correctly
1039    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    /// See [`DecodeStream`]
1067    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
1079/// Internal function exposed only to bypass python limitations
1080pub 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    /// Tokenization logic, makes the bridge between the pre-tokenization phase and the real
1117    /// tokenization phase, and converting offsets back to the original referential.
1118    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    /// Normalization logic, go through all normalizers
1137    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    /// Register the given tokens as special tokens. This is especially useful for removing
1154    /// these special tokens while decoding
1155    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    /// Add the given tokens to the added vocabulary
1161    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    /// PreTokenization logic, handling the case where there is no PreTokenizer set
1172    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    /// Post processing logic, handling the case where there is no PostProcessor set
1190    pub fn post_process(
1191        &self,
1192        encoding: Encoding,
1193        pair_encoding: Option<Encoding>,
1194        add_special_tokens: bool,
1195    ) -> Result<Encoding> {
1196        // 1. First we truncate if needed
1197        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, &params)?
1207                } else {
1208                    truncate_encodings(encoding, pair_encoding, trunc)?
1209                }
1210            } else {
1211                (encoding, pair_encoding)
1212            }
1213        };
1214
1215        // 2. Then We post process
1216        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        // 3. Then we pad if needed
1233        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    /// Encode all the sentences in parallel, using multiple threads
1262    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            // We do the padding here to make sure we handle the batch padding
1277            pad_encodings(&mut encodings, params)?;
1278        }
1279
1280        Ok(encodings)
1281    }
1282
1283    /// Encode all the sentences in parallel, using multiple threads.
1284    /// The offsets on each `Encoding` will be relative to chars instead of bytes.
1285    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            // We do the padding here to make sure we handle the batch padding
1300            pad_encodings(&mut encodings, params)?;
1301        }
1302
1303        Ok(encodings)
1304    }
1305
1306    /// Encode all the sentences in parallel, using multiple threads
1307    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            // We do the padding here to make sure we handle the batch padding
1322            pad_encodings(&mut encodings, params)?;
1323        }
1324
1325        Ok(encodings)
1326    }
1327
1328    /// Decode all sentences in parallel
1329    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    /// Train our Model from files
1344    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                        // We read new lines using this API instead of the Lines Iterator
1363                        // on purpose. We want to keep the `\n` and potential `\r` between each lines
1364                        // We use an iterator to be able to chain with par_bridge.
1365                        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    /// Train our Model, using the given Trainer and iterator
1417    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    /// Instantiate a new Tokenizer from the given file
1491    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    /// Instantiate a new Tokenizer from bytes
1507    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    /// Instantiate a new Tokenizer from a file hosted on the Hugging Face Hub.
1527    /// It expects the `identifier` of a model that includes a `tokenizer.json` file.
1528    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    /// Serialize the current tokenizer as a String
1546    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    /// Save the current tokenizer at the given path
1555    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}