tokenizers/processors/
template.rs

1//! # Template Processing
2//!
3//! Provides a way to specify templates in order to add the special tokens to each
4//! input sequence as relevant.
5//!
6//! ## Example
7//!
8//! Let's take `BERT` tokenizer as an example. It uses two special tokens, used to
9//! delimitate each sequence. `[CLS]` is always used at the beginning of the first
10//! sequence, and `[SEP]` is added at the end of both the first, and the pair
11//! sequences. The final result looks like this:
12//! - Single sequence: `[CLS] Hello there [SEP]`
13//! - Pair sequences: `[CLS] My name is Anthony [SEP] What is my name? [SEP]`
14//!
15//! With the type ids as following:
16//! ```markdown
17//! [CLS]   ...   [SEP]   ...   [SEP]
18//!   0      0      0      1      1
19//! ```
20//!
21//! So, we can define a [`TemplateProcessing`] that will achieve this result:
22//! ```
23//! # use tokenizers::processors::template::TemplateProcessing;
24//! let template = TemplateProcessing::builder()
25//!     // The template when we only have a single sequence:
26//!     .try_single(vec!["[CLS]", "$0", "[SEP]"]).unwrap()
27//!     // Same as:
28//!     .try_single("[CLS] $0 [SEP]").unwrap()
29//!
30//!     // The template when we have both sequences:
31//!     .try_pair(vec!["[CLS]:0", "$A:0", "[SEP]:0", "$B:1", "[SEP]:1"]).unwrap()
32//!     // Same as:
33//!     .try_pair("[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1").unwrap()
34//!     // Or:
35//!     .try_pair("[CLS] $0 [SEP] $B:1 [SEP]:1").unwrap()
36//!
37//!     // The list of special tokens used by each sequences
38//!     .special_tokens(vec![("[CLS]", 1), ("[SEP]", 0)])
39//!     .build()
40//!     .unwrap();
41//! ```
42//!
43//! In this example, each input sequence is identified using a `$` construct. This identifier
44//! lets us specify each input sequence, and the type_id to use. When nothing is specified,
45//! it uses the default values. Here are the different ways to specify it:
46//! - Specifying the sequence, with default `type_id == 0`: `$A` or `$B`
47//! - Specifying the `type_id` with default `sequence == A`: `$0`, `$1`, `$2`, ...
48//! - Specifying both: `$A:0`, `$B:1`, ...
49//!
50//! The same construct is used for special tokens: `<identifier>(:<type_id>)?`.
51//!
52//! **Warning**: You must ensure that you are giving the correct tokens/ids as these will
53//! be added to the `Encoding` without any further check. If the given ids correspond to
54//! something totally different in a `Tokenizer` using this `PostProcessor`, it might lead
55//! to unexpected results.
56//!
57//! [`TemplateProcessing`]: struct.TemplateProcessing.html
58//!
59use crate::{Encoding, PostProcessor, Result};
60use ahash::{AHashMap, AHashSet};
61use itertools::Itertools;
62use serde::{Deserialize, Serialize};
63use std::convert::{TryFrom, TryInto};
64use std::result::Result as StdResult;
65
66/// Represents any sequences received as input of the PostProcessor
67#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq)]
68pub enum Sequence {
69    /// This is the first sequence, the one that is always specified
70    A,
71    /// This is the pair sequence, that is optional
72    B,
73}
74
75/// Represents the different kind of pieces that constitute a template.
76/// It can be either the input sequence or a [`SpecialToken`]:
77///
78/// - The `Sequence` has an associated `type_id` which is used by default
79///   for any token inside this sequence. The `Sequence` corresponds to one
80///   of the input sequence given as input of the `PostProcessor`.
81///
82/// - The `SpecialToken` has an associated `id`. It corresponds to a [`SpecialToken`].
83///
84/// The easiest way to build a `Piece` is actually by converting it from a string:
85/// ```
86/// # use tokenizers::processors::template::Piece;
87/// # use std::convert::TryFrom;
88/// let sequence_with_type_id_0 = Piece::try_from("$0").unwrap();
89/// let sequence_with_type_id_1 = Piece::try_from("$1").unwrap();
90/// let special_token_cls = Piece::try_from("[CLS]").unwrap();
91/// ```
92///
93/// [`SpecialToken`]: struct.SpecialToken.html
94///
95#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq)]
96pub enum Piece {
97    Sequence { id: Sequence, type_id: u32 },
98    SpecialToken { id: String, type_id: u32 },
99}
100
101impl Piece {
102    fn extract_id(s: &str) -> Option<Self> {
103        if s.starts_with('$') {
104            let rest = &s['$'.len_utf8()..];
105
106            // If the id is just `$`, we use 0 as type_id, and Sequence A
107            match rest {
108                "" => Some(Self::Sequence {
109                    id: Sequence::A,
110                    type_id: 0,
111                }),
112                "A" | "a" => Some(Self::Sequence {
113                    id: Sequence::A,
114                    type_id: 0,
115                }),
116                "B" | "b" => Some(Self::Sequence {
117                    id: Sequence::B,
118                    type_id: 0,
119                }),
120                n => {
121                    if let Ok(type_id) = n.parse::<u32>() {
122                        Some(Self::Sequence {
123                            id: Sequence::A,
124                            type_id,
125                        })
126                    } else {
127                        None
128                    }
129                }
130            }
131        } else {
132            Some(Self::SpecialToken {
133                id: s.to_owned(),
134                type_id: 0,
135            })
136        }
137    }
138
139    fn with_type_id(self, type_id: u32) -> Self {
140        match self {
141            Self::Sequence { id, .. } => Self::Sequence { id, type_id },
142            Self::SpecialToken { id, .. } => Self::SpecialToken { id, type_id },
143        }
144    }
145}
146
147impl TryFrom<String> for Piece {
148    type Error = String;
149
150    fn try_from(s: String) -> StdResult<Self, Self::Error> {
151        let parts = s.split(':').collect::<Vec<_>>();
152
153        let err = || format!("Cannot build Piece from string \"{s}\"");
154        match parts.as_slice() {
155            [id, type_id] => {
156                let type_id: u32 = type_id.parse().map_err(|_| err())?;
157                let piece = Self::extract_id(id).ok_or_else(err)?;
158                Ok(piece.with_type_id(type_id))
159            }
160            [id] => Self::extract_id(id).ok_or_else(err),
161            _ => Err(err()),
162        }
163    }
164}
165
166impl TryFrom<&str> for Piece {
167    type Error = String;
168
169    fn try_from(s: &str) -> StdResult<Self, Self::Error> {
170        Piece::try_from(s.to_owned())
171    }
172}
173
174/// Represents a bunch of tokens to be used in a template.
175/// Usually, special tokens have only one associated id/token but in
176/// some cases, it might be interesting to have multiple ids/tokens.
177///
178/// # Examples
179/// ```
180/// # use tokenizers::processors::template::SpecialToken;
181/// // Simple cases, where a single id/token is necessary:
182/// let cls = SpecialToken::from(("[CLS]", 1));
183/// let sep = SpecialToken::from((0, "[SEP]")); // The order in the tuple is not important
184///
185/// // More complex case with multiple values:
186/// let complex = SpecialToken::new(
187///     "A complex special token:".into(),
188///     vec![0, 1, 2, 3, 4],
189///     vec!["A".into(), "complex".into(), "special".into(), "token".into(), ":".into()]
190/// ).unwrap();
191/// ```
192#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq)]
193pub struct SpecialToken {
194    /// A unique id used to identify this SpecialToken in the template
195    id: String,
196    /// The list of associated ids
197    ids: Vec<u32>,
198    /// The list of associated tokens
199    tokens: Vec<String>,
200}
201
202impl From<(String, u32)> for SpecialToken {
203    fn from(v: (String, u32)) -> Self {
204        Self {
205            id: v.0.clone(),
206            ids: vec![v.1],
207            tokens: vec![v.0],
208        }
209    }
210}
211impl From<(&str, u32)> for SpecialToken {
212    fn from(v: (&str, u32)) -> Self {
213        Self::from((v.0.to_owned(), v.1))
214    }
215}
216impl From<(u32, String)> for SpecialToken {
217    fn from(v: (u32, String)) -> Self {
218        Self::from((v.1, v.0))
219    }
220}
221impl From<(u32, &str)> for SpecialToken {
222    fn from(v: (u32, &str)) -> Self {
223        Self::from((v.1.to_owned(), v.0))
224    }
225}
226
227impl SpecialToken {
228    pub fn new(id: String, ids: Vec<u32>, tokens: Vec<String>) -> Result<Self> {
229        if ids.len() != tokens.len() {
230            Err("SpecialToken: ids and tokens must be of the same length".into())
231        } else {
232            Ok(Self { id, ids, tokens })
233        }
234    }
235}
236
237/// A Template represents a Vec<[`Piece`]>.
238///
239/// We can easily build one as follows
240/// ```
241/// # use tokenizers::processors::template::Template;
242/// # use std::convert::TryFrom;
243/// // By providing a `String` or `&str`, we just split on whitespaces:
244/// let template = Template::try_from("[CLS] $0 [SEP]").unwrap();
245///
246/// // By providing pieces directly:
247/// let template = Template::try_from(vec!["[CLS]", "$0", "[SEP]"]).unwrap();
248/// ```
249/// Both of these methods give the same result.
250///
251/// [`Piece`]: enum.Piece.html
252///
253#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq)]
254#[serde(transparent)]
255pub struct Template(Vec<Piece>);
256
257impl<T> TryFrom<Vec<T>> for Template
258where
259    T: TryInto<Piece, Error = String>,
260{
261    type Error = String;
262
263    fn try_from(v: Vec<T>) -> StdResult<Self, Self::Error> {
264        Ok(Self(
265            v.into_iter()
266                .map(|p| p.try_into())
267                .collect::<StdResult<Vec<_>, Self::Error>>()?,
268        ))
269    }
270}
271
272impl TryFrom<String> for Template {
273    type Error = String;
274
275    fn try_from(s: String) -> StdResult<Self, Self::Error> {
276        Self::try_from(s.as_ref())
277    }
278}
279
280impl TryFrom<&str> for Template {
281    type Error = String;
282
283    fn try_from(s: &str) -> StdResult<Self, Self::Error> {
284        Self::try_from(s.split(' ').collect::<Vec<_>>())
285    }
286}
287
288/// A bunch of [`SpecialToken`] represented by their ID.
289/// Internally, `Tokens` is a `HashMap<String, SpecialToken>` and can be built
290/// from a HashMap or a Vec<[`SpecialToken`]>.
291///
292/// [`SpecialToken`]: struct.SpecialToken.html
293#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize, Eq)]
294#[serde(transparent)]
295pub struct Tokens(
296    #[serde(serialize_with = "crate::utils::ordered_map")] pub AHashMap<String, SpecialToken>,
297);
298
299impl<T: Into<SpecialToken>> From<Vec<T>> for Tokens {
300    fn from(v: Vec<T>) -> Self {
301        Self(
302            v.into_iter()
303                .map(|t| {
304                    let token: SpecialToken = t.into();
305                    (token.id.clone(), token)
306                })
307                .collect(),
308        )
309    }
310}
311
312impl From<AHashMap<String, SpecialToken>> for Tokens {
313    fn from(v: AHashMap<String, SpecialToken>) -> Self {
314        Self(v)
315    }
316}
317
318/// This PostProcessor takes care of processing each input `Encoding` by applying
319/// the corresponding template, before merging them in the final Encoding.
320///
321/// A `Template` is actually a sequence of `Piece` that will be
322/// concatenated together in the given order. Each `Piece` represents either
323/// one of the input `Encoding` or a `SpecialToken`.
324///
325/// ## Example
326/// ```
327/// # use tokenizers::processors::template::TemplateProcessing;
328/// let template = TemplateProcessing::builder()
329///     .try_single("[CLS] $A [SEP]").unwrap()
330///     .try_pair("[CLS] $A [SEP] $B:1 [SEP]:1").unwrap()
331///     .special_tokens(vec![("[CLS]", 1), ("[SEP]", 0)])
332///     .build()
333///     .unwrap();
334/// ```
335///
336#[derive(Debug, Clone, PartialEq, Builder, Serialize, Deserialize, Eq)]
337#[serde(tag = "type", from = "TemplateProcessingDeserializer")]
338#[builder(build_fn(validate = "Self::validate"))]
339pub struct TemplateProcessing {
340    #[builder(try_setter, default = "\"$0\".try_into().unwrap()")]
341    pub single: Template,
342    #[builder(try_setter, default = "\"$A:0 $B:1\".try_into().unwrap()")]
343    pair: Template,
344    #[builder(setter(skip), default = "self.default_added(true)")]
345    #[serde(skip)]
346    added_single: usize,
347    #[builder(setter(skip), default = "self.default_added(false)")]
348    #[serde(skip)]
349    added_pair: usize,
350    #[builder(setter(into), default)]
351    special_tokens: Tokens,
352}
353
354impl TemplateProcessing {
355    // Getter for `single`
356    pub fn get_single(&self) -> String {
357        format!("{:?}", self.single)
358    }
359
360    // Setter for `single`
361    pub fn set_single(&mut self, single: Template) {
362        self.single = single;
363    }
364
365    // Getter for `pair`
366    pub fn get_pair(&self) -> &Template {
367        &self.pair
368    }
369
370    // Setter for `pair`
371    pub fn set_pair(&mut self, pair: Template) {
372        self.pair = pair;
373    }
374
375    // Getter for `added_single`
376    pub fn get_added_single(&self) -> usize {
377        self.added_single
378    }
379
380    // Setter for `added_single`
381    pub fn set_added_single(&mut self, added_single: usize) {
382        self.added_single = added_single;
383    }
384
385    // Getter for `added_pair`
386    pub fn get_added_pair(&self) -> usize {
387        self.added_pair
388    }
389
390    // Setter for `added_pair`
391    pub fn set_added_pair(&mut self, added_pair: usize) {
392        self.added_pair = added_pair;
393    }
394
395    // Getter for `special_tokens`
396    pub fn get_special_tokens(&self) -> &Tokens {
397        &self.special_tokens
398    }
399
400    // Setter for `special_tokens`
401    pub fn set_special_tokens(&mut self, special_tokens: Tokens) {
402        self.special_tokens = special_tokens;
403    }
404}
405
406impl From<&str> for TemplateProcessingBuilderError {
407    fn from(e: &str) -> Self {
408        e.to_string().into()
409    }
410}
411
412impl PartialEq for TemplateProcessingBuilderError {
413    fn eq(&self, other: &Self) -> bool {
414        self.to_string() == other.to_string()
415    }
416}
417
418/// We use this custom deserializer to provided the values for `added_single`
419/// and `added_pair` during deserialization, while not having to serialize them
420#[doc(hidden)]
421#[derive(Deserialize)]
422#[serde(tag = "type")]
423struct TemplateProcessingDeserializer {
424    single: Template,
425    pair: Template,
426    special_tokens: Tokens,
427}
428impl From<TemplateProcessingDeserializer> for TemplateProcessing {
429    fn from(t: TemplateProcessingDeserializer) -> Self {
430        let added_single = count_added(&t.single, Some(&t.special_tokens));
431        let added_pair = count_added(&t.pair, Some(&t.special_tokens));
432        Self {
433            single: t.single,
434            pair: t.pair,
435            added_single,
436            added_pair,
437            special_tokens: t.special_tokens,
438        }
439    }
440}
441
442/// Count the number of added tokens in the given template
443fn count_added(container: &Template, special_tokens: Option<&Tokens>) -> usize {
444    container
445        .0
446        .iter()
447        .map(|p| match p {
448            Piece::Sequence { .. } => 0,
449            Piece::SpecialToken { id, .. } => {
450                special_tokens.map_or(0, |spt| spt.0.get(id).map_or(0, |s| s.ids.len()))
451            }
452        })
453        .sum()
454}
455
456impl TemplateProcessingBuilder {
457    fn default_added(&self, is_single: bool) -> usize {
458        let container = if is_single {
459            self.single.as_ref()
460        } else {
461            self.pair.as_ref()
462        };
463        container.map_or(0, |pieces| {
464            count_added(pieces, self.special_tokens.as_ref())
465        })
466    }
467
468    fn validate(&self) -> std::result::Result<(), String> {
469        let pair_has_both = self.pair.as_ref().is_none_or(|pair| {
470            let mut has_a = false;
471            let mut has_b = false;
472            for piece in &pair.0 {
473                if let Piece::Sequence {
474                    id: Sequence::A, ..
475                } = piece
476                {
477                    has_a = true;
478                }
479                if let Piece::Sequence {
480                    id: Sequence::B, ..
481                } = piece
482                {
483                    has_b = true;
484                }
485            }
486            has_a && has_b
487        });
488        if !pair_has_both {
489            return Err("Template for `pair` must use both sequences".into());
490        }
491
492        let check = |sp| {
493            let exist = self
494                .special_tokens
495                .as_ref()
496                .is_some_and(|map| map.0.contains_key(sp));
497
498            match exist {
499                false => Some(sp),
500                true => None,
501            }
502        };
503
504        let empty = [];
505        let missing: AHashSet<&str> = self
506            .single
507            .as_ref()
508            .map_or(empty.iter(), |s| s.0.iter())
509            .chain(self.pair.as_ref().map_or(empty.iter(), |s| s.0.iter()))
510            .filter_map(|piece| match piece {
511                Piece::Sequence { .. } => None,
512                Piece::SpecialToken { id, .. } => check(id.as_ref()),
513            })
514            .collect::<AHashSet<_>>();
515
516        if missing.is_empty() {
517            Ok(())
518        } else {
519            Err(format!(
520                "Missing SpecialToken(s) with id(s) `{}`",
521                missing.iter().join(", ")
522            ))
523        }
524    }
525}
526
527impl Default for TemplateProcessing {
528    fn default() -> Self {
529        Self {
530            single: "$0".try_into().unwrap(),
531            pair: "$1".try_into().unwrap(),
532            added_single: 0,
533            added_pair: 0,
534            special_tokens: Tokens::default(),
535        }
536    }
537}
538
539impl TemplateProcessing {
540    pub fn builder() -> TemplateProcessingBuilder {
541        TemplateProcessingBuilder::default()
542    }
543
544    fn apply_template(
545        &self,
546        template: &[Piece],
547        mut encodings: Vec<Encoding>,
548        add_special_tokens: bool,
549    ) -> Result<Vec<Encoding>> {
550        let final_encodings: Vec<Encoding> = template
551            .iter()
552            .flat_map(|piece| {
553                match piece {
554                    Piece::Sequence { id, type_id } => {
555                        let i = usize::from(*id != Sequence::A);
556                        let encoding = &mut encodings[i];
557                        encoding.set_type_ids(vec![*type_id; encoding.len()]);
558                        encoding.set_sequence_id(i);
559                        Some(encoding.clone())
560                    }
561                    Piece::SpecialToken { id, type_id } => {
562                        if add_special_tokens {
563                            let tok = &self.special_tokens.0[id]; // We already checked existence above
564                            let len = tok.ids.len();
565
566                            let encoding = Encoding::new(
567                                tok.ids.clone(),
568                                std::iter::repeat_n(*type_id, len).collect(),
569                                tok.tokens.clone(),
570                                // words
571                                std::iter::repeat_n(None, len).collect(),
572                                // offsets
573                                std::iter::repeat_n((0, 0), len).collect(),
574                                // special_tokens_mask
575                                std::iter::repeat_n(1, len).collect(),
576                                // attention_mask
577                                std::iter::repeat_n(1, len).collect(),
578                                // overflowing
579                                vec![],
580                                // sequence_range
581                                AHashMap::new(),
582                            );
583                            Some(encoding)
584                        } else {
585                            None
586                        }
587                    }
588                }
589            })
590            .collect();
591
592        //let mut pair = if encodings.len() > 1 {
593        //    Some(encodings.pop().unwrap())
594        //} else {
595        //    None
596        //};
597        //let mut encoding = encodings.pop().unwrap();
598
599        //let pair_overflowing = pair.as_mut().map_or(vec![], |e| e.take_overflowing());
600        //let mut overflowing: Vec<Encoding> = encoding
601        //    .take_overflowing()
602        //    .iter()
603        //    .map(|encoding| -> Result<Vec<Encoding>> {
604        //        // 1. The pair itself
605        //        let mut overflowings = self.apply_template(
606        //            template,
607        //            if encodings.len() > 1 {
608        //                vec![encoding.clone(), encodings[1].clone()]
609        //            } else {
610        //                vec![encoding.clone()]
611        //            },
612        //            add_special_tokens,
613        //        )?;
614
615        //        // 2. Its overflowings
616        //        for other_o in &pair_overflowing {
617        //            overflowings.extend(self.apply_template(
618        //                template,
619        //                vec![encoding.clone(), other_o.clone()],
620        //                add_special_tokens,
621        //            )?);
622        //        }
623
624        //        Ok(overflowings)
625        //    })
626        //    .collect::<Result<Vec<Vec<Encoding>>>>()?
627        //    .into_iter()
628        //    .flatten()
629        //    .collect();
630        //// We also need to combine the first sequence with all other overflowings
631        //overflowing.extend(
632        //    pair_overflowing
633        //        .into_iter()
634        //        .map(|pair| {
635        //            self.apply_template(template, vec![encoding.clone(), pair], add_special_tokens)
636        //        })
637        //        .collect::<Result<Vec<_>>>()?
638        //        .into_iter()
639        //        .flatten(),
640        //);
641
642        Ok(final_encodings)
643    }
644}
645
646impl PostProcessor for TemplateProcessing {
647    fn added_tokens(&self, is_pair: bool) -> usize {
648        if is_pair {
649            self.added_pair
650        } else {
651            self.added_single
652        }
653    }
654
655    fn process_encodings(
656        &self,
657        encodings: Vec<Encoding>,
658        add_special_tokens: bool,
659    ) -> Result<Vec<Encoding>> {
660        // let (encoding, pair): (Encoding, Option<Encoding>) = match encodings.len() {
661        //     1 => (
662        //         encodings
663        //             .pop()
664        //             .ok_or(ProcessorError::InvalidEncodingsVecLength)?,
665        //         None,
666        //     ),
667        //     2 => {
668        //         let pair = encodings
669        //             .pop()
670        //             .ok_or(ProcessorError::InvalidEncodingsVecLength)?;
671        //         let encoding = encodings
672        //             .pop()
673        //             .ok_or(ProcessorError::InvalidEncodingsVecLength)?;
674        //         (encoding, Some(pair))
675        //     }
676        //     _ => return Err(Box::new(ProcessorError::InvalidEncodingsVecLength)),
677        // };
678        let template = match encodings.len() {
679            2 => &self.pair.0,
680            1 => &self.single.0,
681            _ => todo!(),
682        };
683        let encodings = self.apply_template(template, encodings, add_special_tokens)?;
684        Ok(encodings)
685    }
686}
687
688#[cfg(test)]
689mod tests {
690    use super::*;
691    use std::convert::TryInto;
692    use std::iter::FromIterator;
693
694    #[test]
695    fn piece_serde() {
696        let seq_0 = Piece::Sequence {
697            id: Sequence::A,
698            type_id: 0,
699        };
700        let seq_0_s = r#"{"Sequence":{"id":"A","type_id":0}}"#;
701
702        assert_eq!(serde_json::to_string(&seq_0).unwrap(), seq_0_s);
703        assert_eq!(serde_json::from_str::<Piece>(seq_0_s).unwrap(), seq_0);
704
705        let seq_1 = Piece::Sequence {
706            id: Sequence::B,
707            type_id: 1,
708        };
709        let seq_1_s = r#"{"Sequence":{"id":"B","type_id":1}}"#;
710        assert_eq!(serde_json::to_string(&seq_1).unwrap(), seq_1_s);
711        assert_eq!(serde_json::from_str::<Piece>(seq_1_s).unwrap(), seq_1);
712
713        let spe = Piece::SpecialToken {
714            id: "[CLS]".into(),
715            type_id: 0,
716        };
717        let spe_s = r#"{"SpecialToken":{"id":"[CLS]","type_id":0}}"#;
718        assert_eq!(serde_json::to_string(&spe).unwrap(), spe_s);
719        assert_eq!(serde_json::from_str::<Piece>(spe_s).unwrap(), spe);
720    }
721
722    #[test]
723    fn piece() {
724        assert_eq!(
725            Ok(Piece::Sequence {
726                id: Sequence::A,
727                type_id: 0
728            }),
729            "$".try_into()
730        );
731        assert_eq!(
732            Ok(Piece::Sequence {
733                id: Sequence::B,
734                type_id: 0
735            }),
736            "$B".try_into()
737        );
738        assert_eq!(
739            Ok(Piece::Sequence {
740                id: Sequence::A,
741                type_id: 1
742            }),
743            "$1".try_into()
744        );
745        assert_eq!(
746            Ok(Piece::Sequence {
747                id: Sequence::B,
748                type_id: 2
749            }),
750            "$B:2".try_into()
751        );
752        assert_eq!(
753            Ok(Piece::Sequence {
754                id: Sequence::A,
755                type_id: 1
756            }),
757            "$:1".try_into()
758        );
759        assert!(Piece::try_from("$C:1").is_err());
760        assert!(Piece::try_from("$A:").is_err());
761    }
762
763    #[test]
764    fn special_token_serde() {
765        let simple = SpecialToken::from(("[CLS]", 0));
766        let simple_s = r#"{"id":"[CLS]","ids":[0],"tokens":["[CLS]"]}"#;
767        assert_eq!(serde_json::to_string(&simple).unwrap(), simple_s);
768        assert_eq!(
769            serde_json::from_str::<SpecialToken>(simple_s).unwrap(),
770            simple
771        );
772
773        let complete = SpecialToken::new(
774            "[2FR]".into(),
775            vec![1, 2, 3],
776            vec!["convert".into(), "to".into(), "FR".into()],
777        )
778        .unwrap();
779        let complete_s = r#"{"id":"[2FR]","ids":[1,2,3],"tokens":["convert","to","FR"]}"#;
780        assert_eq!(serde_json::to_string(&complete).unwrap(), complete_s);
781        assert_eq!(
782            serde_json::from_str::<SpecialToken>(complete_s).unwrap(),
783            complete
784        );
785
786        let malformed = SpecialToken::new(
787            "[2FR]".into(),
788            vec![1, 2],
789            vec!["convert".into(), "to".into(), "FR".into()],
790        );
791        assert!(malformed.is_err());
792        let malformed = SpecialToken::new(
793            "[2FR]".into(),
794            vec![1, 2, 3],
795            vec!["convert".into(), "FR".into()],
796        );
797        assert!(malformed.is_err());
798    }
799
800    #[test]
801    fn template_serde() {
802        let template = Template(vec![
803            Piece::Sequence {
804                id: Sequence::A,
805                type_id: 0,
806            },
807            Piece::SpecialToken {
808                id: "[CLS]".into(),
809                type_id: 0,
810            },
811        ]);
812        let template_s =
813            r#"[{"Sequence":{"id":"A","type_id":0}},{"SpecialToken":{"id":"[CLS]","type_id":0}}]"#;
814        assert_eq!(serde_json::to_string(&template).unwrap(), template_s);
815        assert_eq!(
816            serde_json::from_str::<Template>(template_s).unwrap(),
817            template
818        );
819    }
820
821    #[test]
822    fn tokens_serde() {
823        let tokens = Tokens::from(vec![("[CLS]", 1), ("[SEP]", 0)]);
824        let tokens_s = r#"{"[CLS]":{"id":"[CLS]","ids":[1],"tokens":["[CLS]"]},"[SEP]":{"id":"[SEP]","ids":[0],"tokens":["[SEP]"]}}"#;
825        let tokens_ser = serde_json::to_string(&tokens).unwrap();
826        assert_eq!(tokens_ser, tokens_s);
827        assert_eq!(serde_json::from_str::<Tokens>(tokens_s).unwrap(), tokens);
828    }
829
830    fn get_bert_template() -> TemplateProcessing {
831        TemplateProcessing::builder()
832            .try_single(vec!["[CLS]", "$0", "[SEP]"])
833            .unwrap()
834            .try_pair("[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1")
835            .unwrap()
836            .special_tokens(vec![("[CLS]", 1), ("[SEP]", 0)])
837            .build()
838            .unwrap()
839    }
840
841    #[test]
842    fn template_processing_serde() {
843        let template = tests::get_bert_template();
844        let template_s = "{\
845            \"type\":\"TemplateProcessing\",\
846            \"single\":[\
847                {\"SpecialToken\":{\"id\":\"[CLS]\",\"type_id\":0}},\
848                {\"Sequence\":{\"id\":\"A\",\"type_id\":0}},\
849                {\"SpecialToken\":{\"id\":\"[SEP]\",\"type_id\":0}}\
850            ],\
851            \"pair\":[\
852                {\"SpecialToken\":{\"id\":\"[CLS]\",\"type_id\":0}},\
853                {\"Sequence\":{\"id\":\"A\",\"type_id\":0}},\
854                {\"SpecialToken\":{\"id\":\"[SEP]\",\"type_id\":0}},\
855                {\"Sequence\":{\"id\":\"B\",\"type_id\":1}},\
856                {\"SpecialToken\":{\"id\":\"[SEP]\",\"type_id\":1}}\
857            ],\
858            \"special_tokens\":{\
859                \"[CLS]\":{\
860                    \"id\":\"[CLS]\",\"ids\":[1],\"tokens\":[\"[CLS]\"]\
861                },\
862                \"[SEP]\":{\
863                    \"id\":\"[SEP]\",\"ids\":[0],\"tokens\":[\"[SEP]\"]\
864                }\
865            }}";
866        let template_ser = serde_json::to_string(&template).unwrap();
867        assert_eq!(template_ser, template_s);
868        assert_eq!(
869            serde_json::from_str::<TemplateProcessing>(template_s).unwrap(),
870            template
871        );
872    }
873
874    #[test]
875    fn missing_special_tokens() {
876        let processor = TemplateProcessing::builder()
877            .try_single("[CLS] $0 [SEP]")
878            .unwrap()
879            .try_pair("[CLS] $A:0 [SEP] $B:1 [SEP]")
880            .unwrap()
881            .build();
882
883        let err_a = Err("Missing SpecialToken(s) with id(s) `[SEP], [CLS]`".into());
884        let err_b = Err("Missing SpecialToken(s) with id(s) `[CLS], [SEP]`".into());
885        assert!(processor == err_a || processor == err_b);
886    }
887
888    #[test]
889    fn template_processing() {
890        let processor = tests::get_bert_template();
891        assert_eq!(processor.added_tokens(false), 2);
892        assert_eq!(processor.added_tokens(true), 3);
893
894        use crate::Token;
895        let encoding = Encoding::from_tokens(
896            vec![
897                Token::new(12, "Hello".into(), (0, 5)),
898                Token::new(14, "there".into(), (6, 11)),
899            ],
900            0,
901        );
902        let pair = Encoding::from_tokens(vec![Token::new(15, "pair".into(), (0, 4))], 0);
903        let single_encoding = processor.process(encoding.clone(), None, true).unwrap();
904        assert_eq!(
905            single_encoding,
906            Encoding::new(
907                vec![1, 12, 14, 0],
908                vec![0, 0, 0, 0],
909                vec![
910                    "[CLS]".into(),
911                    "Hello".into(),
912                    "there".into(),
913                    "[SEP]".into()
914                ],
915                vec![None, None, None, None],
916                vec![(0, 0), (0, 5), (6, 11), (0, 0)],
917                vec![1, 0, 0, 1],
918                vec![1, 1, 1, 1],
919                vec![],
920                AHashMap::from_iter(vec![(0, 1..3)]),
921            )
922        );
923        assert_eq!(single_encoding.token_to_sequence(2), Some(0));
924        assert_eq!(single_encoding.token_to_sequence(3), None);
925        let pair_encoding = processor.process(encoding, Some(pair), true).unwrap();
926        assert_eq!(
927            pair_encoding,
928            Encoding::new(
929                vec![1, 12, 14, 0, 15, 0],
930                vec![0, 0, 0, 0, 1, 1],
931                vec![
932                    "[CLS]".into(),
933                    "Hello".into(),
934                    "there".into(),
935                    "[SEP]".into(),
936                    "pair".into(),
937                    "[SEP]".into()
938                ],
939                vec![None, None, None, None, None, None],
940                vec![(0, 0), (0, 5), (6, 11), (0, 0), (0, 4), (0, 0)],
941                vec![1, 0, 0, 1, 0, 1],
942                vec![1, 1, 1, 1, 1, 1],
943                vec![],
944                AHashMap::from_iter(vec![(0, 1..3), (1, 4..5)]),
945            )
946        );
947        assert_eq!(pair_encoding.token_to_sequence(2), Some(0));
948        assert_eq!(pair_encoding.token_to_sequence(3), None);
949        assert_eq!(pair_encoding.token_to_sequence(4), Some(1));
950        assert_eq!(pair_encoding.token_to_sequence(5), None);
951    }
952
953    #[test]
954    fn template_processing_overflowing() {
955        let processor = tests::get_bert_template();
956        assert_eq!(processor.added_tokens(false), 2);
957        assert_eq!(processor.added_tokens(true), 3);
958
959        use crate::Token;
960        let mut encoding = Encoding::from_tokens(
961            vec![
962                Token::new(12, "Hello".into(), (0, 5)),
963                Token::new(14, "there".into(), (6, 11)),
964            ],
965            0,
966        );
967        let overflowing = Encoding::from_tokens(vec![Token::new(13, "you".into(), (12, 15))], 0);
968        encoding.set_overflowing(vec![overflowing]);
969
970        let mut pair = Encoding::from_tokens(
971            vec![
972                Token::new(15, "pair".into(), (0, 4)),
973                Token::new(16, "with".into(), (5, 9)),
974            ],
975            0,
976        );
977        let pair_overflowing =
978            Encoding::from_tokens(vec![Token::new(17, "info".into(), (10, 14))], 0);
979        pair.set_overflowing(vec![pair_overflowing]);
980
981        let single_encoding = processor.process(encoding.clone(), None, true).unwrap();
982        assert_eq!(
983            single_encoding,
984            Encoding::new(
985                vec![1, 12, 14, 0],
986                vec![0, 0, 0, 0],
987                vec![
988                    "[CLS]".into(),
989                    "Hello".into(),
990                    "there".into(),
991                    "[SEP]".into()
992                ],
993                vec![None, None, None, None],
994                vec![(0, 0), (0, 5), (6, 11), (0, 0)],
995                vec![1, 0, 0, 1],
996                vec![1, 1, 1, 1],
997                vec![Encoding::new(
998                    vec![1, 13, 0],
999                    vec![0, 0, 0],
1000                    vec!["[CLS]".into(), "you".into(), "[SEP]".into()],
1001                    vec![None, None, None],
1002                    vec![(0, 0), (12, 15), (0, 0)],
1003                    vec![1, 0, 1],
1004                    vec![1, 1, 1],
1005                    vec![],
1006                    AHashMap::from_iter(vec![(0, 1..2)]),
1007                )],
1008                AHashMap::from_iter(vec![(0, 1..3)]),
1009            )
1010        );
1011        assert_eq!(single_encoding.token_to_sequence(2), Some(0));
1012        assert_eq!(single_encoding.token_to_sequence(3), None);
1013        let pair_encoding = processor.process(encoding, Some(pair), true).unwrap();
1014        println!("{pair_encoding:#?}");
1015        assert_eq!(
1016            pair_encoding,
1017            Encoding::new(
1018                vec![1, 12, 14, 0, 15, 16, 0],
1019                vec![0, 0, 0, 0, 1, 1, 1],
1020                vec![
1021                    "[CLS]".into(),
1022                    "Hello".into(),
1023                    "there".into(),
1024                    "[SEP]".into(),
1025                    "pair".into(),
1026                    "with".into(),
1027                    "[SEP]".into()
1028                ],
1029                vec![None, None, None, None, None, None, None],
1030                vec![(0, 0), (0, 5), (6, 11), (0, 0), (0, 4), (5, 9), (0, 0)],
1031                vec![1, 0, 0, 1, 0, 0, 1],
1032                vec![1, 1, 1, 1, 1, 1, 1],
1033                vec![
1034                    Encoding::new(
1035                        vec![1, 13, 0, 15, 16, 0],
1036                        vec![0, 0, 0, 1, 1, 1],
1037                        vec![
1038                            "[CLS]".into(),
1039                            "you".into(),
1040                            "[SEP]".into(),
1041                            "pair".into(),
1042                            "with".into(),
1043                            "[SEP]".into()
1044                        ],
1045                        vec![None, None, None, None, None, None],
1046                        vec![(0, 0), (12, 15), (0, 0), (0, 4), (5, 9), (0, 0)],
1047                        vec![1, 0, 1, 0, 0, 1],
1048                        vec![1, 1, 1, 1, 1, 1],
1049                        vec![Encoding::new(
1050                            vec![1, 13, 0, 17, 0],
1051                            vec![0, 0, 0, 0, 1],
1052                            vec![
1053                                "[CLS]".into(),
1054                                "you".into(),
1055                                "[SEP]".into(),
1056                                "info".into(),
1057                                "[SEP]".into()
1058                            ],
1059                            vec![None, None, None, None, None,],
1060                            vec![(0, 0), (12, 15), (0, 0), (10, 14), (0, 0)],
1061                            vec![1, 0, 1, 0, 1],
1062                            vec![1, 1, 1, 1, 1],
1063                            vec![],
1064                            AHashMap::from_iter(vec![(0, 1..2), (1, 3..4)]),
1065                        ),],
1066                        AHashMap::from_iter(vec![(1, 3..5), (0, 1..2)]),
1067                    ),
1068                    Encoding::new(
1069                        vec![1, 13, 0, 17, 0],
1070                        vec![0, 0, 0, 0, 1],
1071                        vec![
1072                            "[CLS]".into(),
1073                            "you".into(),
1074                            "[SEP]".into(),
1075                            "info".into(),
1076                            "[SEP]".into()
1077                        ],
1078                        vec![None, None, None, None, None,],
1079                        vec![(0, 0), (12, 15), (0, 0), (10, 14), (0, 0)],
1080                        vec![1, 0, 1, 0, 1],
1081                        vec![1, 1, 1, 1, 1],
1082                        vec![],
1083                        AHashMap::from_iter(vec![(0, 1..2), (1, 3..4)]),
1084                    ),
1085                    Encoding::new(
1086                        vec![1, 12, 14, 0, 17, 0],
1087                        vec![0, 0, 0, 0, 0, 1],
1088                        vec![
1089                            "[CLS]".into(),
1090                            "Hello".into(),
1091                            "there".into(),
1092                            "[SEP]".into(),
1093                            "info".into(),
1094                            "[SEP]".into()
1095                        ],
1096                        vec![None, None, None, None, None, None],
1097                        vec![(0, 0), (0, 5), (6, 11), (0, 0), (10, 14), (0, 0)],
1098                        vec![1, 0, 0, 1, 0, 1],
1099                        vec![1, 1, 1, 1, 1, 1],
1100                        vec![Encoding::new(
1101                            vec![1, 13, 0, 17, 0],
1102                            vec![0, 0, 0, 0, 1],
1103                            vec![
1104                                "[CLS]".into(),
1105                                "you".into(),
1106                                "[SEP]".into(),
1107                                "info".into(),
1108                                "[SEP]".into()
1109                            ],
1110                            vec![None, None, None, None, None,],
1111                            vec![(0, 0), (12, 15), (0, 0), (10, 14), (0, 0)],
1112                            vec![1, 0, 1, 0, 1],
1113                            vec![1, 1, 1, 1, 1],
1114                            vec![],
1115                            AHashMap::from_iter(vec![(0, 1..2), (1, 3..4)]),
1116                        ),],
1117                        AHashMap::from_iter(vec![(0, 1..3), (1, 4..5)]),
1118                    )
1119                ],
1120                AHashMap::from_iter(vec![(0, 1..3), (1, 4..6)]),
1121            )
1122        );
1123        assert_eq!(pair_encoding.token_to_sequence(2), Some(0));
1124        assert_eq!(pair_encoding.token_to_sequence(3), None);
1125        assert_eq!(pair_encoding.token_to_sequence(4), Some(1));
1126        assert_eq!(pair_encoding.token_to_sequence(5), Some(1));
1127        assert_eq!(pair_encoding.token_to_sequence(6), None);
1128    }
1129    #[test]
1130    fn pair_must_use_both_sequences() {
1131        let processor = TemplateProcessing::builder()
1132            .try_single("$0")
1133            .unwrap()
1134            .try_pair("$0 $1")
1135            .unwrap()
1136            .build();
1137        assert_eq!(
1138            processor,
1139            Err("Template for `pair` must use both sequences".into())
1140        );
1141    }
1142
1143    #[test]
1144    fn expect_wrong_error_message() {
1145        let processor = TemplateProcessing::builder()
1146            .try_single("$0")
1147            .unwrap()
1148            .try_pair("$0 $1")
1149            .unwrap()
1150            .build();
1151        assert_ne!(
1152            processor,
1153            Err("Expect the left side error message to be different from the right side!".into())
1154        );
1155    }
1156}
tokenizers/processors/template.rs

tokenizers/processors/
template.rs
