Ability to encode with added tokens

tokenizers/src/tokenizer/mod.rs

@@ -191,6 +191,27 @@ impl Tokenizer {
     /// Encode the given sentence
     pub fn encode(&self, input: EncodeInput) -> Result<Encoding> {
         let generate_output = move |sentence: String, type_id: u32| -> Result<Encoding> {
+            // First we need to split into as many sequences as needed to avoid splitting
+            // on our added tokens
+            let mut encodings = self
+                .split_on_added_tokens(&sentence)
+                .into_iter()
+                .map(|(sentence, id)| -> Result<Encoding> {
+                    // If this is one of our added tokens, lets return an encoding directly
+                    if let Some(id) = id {
+                        return Ok(Encoding::new(
+                            sentence.clone(),
+                            sentence.clone(),
+                            vec![id],
+                            vec![type_id],
+                            vec![sentence.to_owned()],
+                            vec![(0, sentence.len())],
+                            vec![0],
+                            vec![1],
+                            None,
+                        ));
+                    }
+
                     // 1. Normalization
                     // TODO: Make sure we have the offsets update necessary to go from the original text to
                     // the normalized one
@@ -229,6 +250,21 @@ impl Tokenizer {
                         vec![1; length],
                         None,
                     ))
+                })
+                .collect::<Result<Vec<Encoding>>>()?;
+
+            if encodings.is_empty() {
+                return Ok(Encoding::default());
+            }
+
+            let others = encodings.split_off(1);
+            let mut first: Encoding = encodings.into_iter().nth(0).unwrap();
+
+            for encoding in others {
+                first.merge_with(encoding);
+            }
+
+            Ok(first)
         };
 
         let (sentence, pair) = match input {
@@ -412,4 +448,54 @@ impl Tokenizer {
         // Return the number of added tokens
         tokens.len() - ignored
     }
+
+    /// Split the given sentence on multiple parts, finding the added tokens and their id in the process
+    fn split_on_added_tokens(&self, sentence: &str) -> Vec<(String, Option<u32>)> {
+        if let Some(split_re) = &self.split_re {
+            let splits = split_re
+                .find_iter(&sentence)
+                .map(|m| (m.start(), m.end()))
+                .collect::<Vec<_>>();
+
+            // We also insert the splits that are inbetween the added tokens, to split the entire string
+            let mut start_offset = 0;
+            let mut splits = splits
+                .into_iter()
+                .map(|(start, end)| {
+                    let mut splits = vec![];
+                    if start_offset < start {
+                        splits.push((start_offset, start));
+                    }
+                    splits.push((start, end));
+                    start_offset = end;
+
+                    splits
+                })
+                .flatten()
+                .collect::<Vec<_>>();
+            if let Some((_, end)) = splits.iter().last().copied() {
+                if end < sentence.len() {
+                    splits.push((end, sentence.len()));
+                }
+            }
+
+            if splits.is_empty() {
+                vec![(sentence.to_owned(), None)]
+            } else {
+                splits
+                    .into_iter()
+                    .map(|(start, end)| unsafe {
+                        let s = sentence.get_unchecked(start..end).to_owned();
+                        let id = self.added_tokens.get(&AddedToken {
+                            content: s.clone(),
+                            ..Default::default()
+                        });
+                        (s, id.copied())
+                    })
+                    .collect()
+            }
+        } else {
+            vec![(sentence.to_owned(), None)]
+        }
+    }
 }