[pre_tokenizers] Fix sentencepiece based Metaspace

bindings/python/src/pre_tokenizers.rs

@@ -489,6 +489,16 @@ impl PyMetaspace {
         setter!(self_, Metaspace, add_prefix_space, add_prefix_space);
     }
 
+    #[getter]
+    fn get_legacy(self_: PyRef<Self>) -> bool {
+        getter!(self_, Metaspace, legacy)
+    }
+
+    #[setter]
+    fn set_legacy(self_: PyRef<Self>, legacy: bool) {
+        setter!(self_, Metaspace, legacy, legacy);
+    }
+
     #[new]
     #[pyo3(signature = (replacement = PyChar('▁'), add_prefix_space = true, **_kwargs), text_signature = "(self, replacement=\"_\", add_prefix_space=True)")]
     fn new(

bindings/python/tests/bindings/test_pre_tokenizers.py

@@ -110,6 +110,8 @@ def test_can_modify(self):
         assert pretok.replacement == "%"
         pretok.add_prefix_space = True
         assert pretok.add_prefix_space == True
+        pretok.legacy = False
+        assert pretok.legacy == False
 
 
 class TestCharDelimiterSplit:

tokenizers/src/decoders/mod.rs

@@ -78,7 +78,6 @@ mod tests {
         let serialized = serde_json::to_string(&decoder).unwrap();
         assert_eq!(serialized, json);
     }
-
     #[test]
     fn decoder_serialization_other_no_arg() {
         let json = r#"{"type":"Sequence","decoders":[{"type":"Fuse"},{"type":"Metaspace","replacement":"▁","add_prefix_space":true}]}"#;

tokenizers/src/pre_tokenizers/metaspace.rs

@@ -9,10 +9,20 @@ use crate::tokenizer::{Decoder, PreTokenizedString, PreTokenizer, Result, SplitD
 pub struct Metaspace {
     replacement: char,
     pub add_prefix_space: bool,
+    #[serde(skip_serializing_if = "skip_legacy_serialization")]
+    pub legacy: bool,
     #[serde(skip)]
     str_rep: String,
 }
 
+fn skip_legacy_serialization(legacy: &bool) -> bool {
+    *legacy // Skip serialization if legacy is true
+}
+
+fn default_legacy_value() -> bool {
+    true
+}
+
 impl<'de> Deserialize<'de> for Metaspace {
     fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
     where
@@ -29,12 +39,16 @@ impl<'de> Deserialize<'de> for Metaspace {
             _type: Type,
             replacement: char,
             pub add_prefix_space: bool,
+            #[serde(default = "default_legacy_value")]
+            pub legacy: bool,
             #[serde(skip, rename = "str_rep")]
             _str_rep: String,
         }
 
         let helper = MetaspaceHelper::deserialize(deserializer)?;
-        Ok(Self::new(helper.replacement, helper.add_prefix_space))
+        let mut instance = Self::new(helper.replacement, helper.add_prefix_space);
+        instance.legacy = helper.legacy;
+        Ok(instance)
     }
 }
 
@@ -44,6 +58,7 @@ impl Metaspace {
             replacement,
             str_rep: replacement.to_string(),
             add_prefix_space,
+            legacy: true,
         }
     }
 
@@ -65,12 +80,18 @@ impl Default for Metaspace {
 
 impl PreTokenizer for Metaspace {
     fn pre_tokenize(&self, pretokenized: &mut PreTokenizedString) -> Result<()> {
+        let mut first_split = true;
+
         pretokenized.split(|_, mut normalized| {
             normalized.replace(' ', &self.str_rep)?;
             if self.add_prefix_space && !normalized.get().starts_with(self.replacement) {
-                normalized.prepend(&self.str_rep);
+                if self.legacy {
+                    normalized.prepend(&self.str_rep);
+                } else if first_split {
+                    normalized.prepend(&self.str_rep);
+                    first_split = false; // Set the flag to false after the first split
+                }
             }
-
             normalized.split(self.replacement, SplitDelimiterBehavior::MergedWithNext)
         })
     }
@@ -103,6 +124,8 @@ impl Decoder for Metaspace {
 
 #[cfg(test)]
 mod tests {
+    use regex::Regex;
+
     use super::*;
     use crate::{OffsetReferential, OffsetType};
 
@@ -188,6 +211,55 @@ mod tests {
         );
     }
 
+    #[test]
+    fn non_legacy_meta_space() {
+        let mut pretok = Metaspace::new('▁', true);
+        pretok.legacy = false;
+        let mut pretokenized = PreTokenizedString::from("Hey my friend <s>how▁are you");
+        let re_ref = Regex::new(r"(<s>)").unwrap();
+        pretokenized
+            .split(|_, sequence| sequence.split(&re_ref, SplitDelimiterBehavior::Isolated))
+            .expect("Bad split");
+        println!("{:?}", pretokenized);
+
+        pretok.pre_tokenize(&mut pretokenized).unwrap();
+        assert_eq!(
+            pretokenized
+                .get_splits(OffsetReferential::Normalized, OffsetType::Byte)
+                .into_iter()
+                .map(|(s, o, _)| (s, o))
+                .collect::<Vec<_>>(),
+            vec![
+                ("▁Hey", (0, 6)),
+                ("▁my", (6, 11)),
+                ("▁friend", (11, 20)),
+                ("▁", (20, 23)),
+                ("<s>", (23, 26)),
+                ("how", (26, 29)),
+                ("▁are", (29, 35)),
+                ("▁you", (35, 41))
+            ]
+        );
+        pretok.legacy = true;
+        pretok.pre_tokenize(&mut pretokenized).unwrap();
+        assert_eq!(
+            pretokenized
+                .get_splits(OffsetReferential::Normalized, OffsetType::Byte)
+                .into_iter()
+                .map(|(s, o, _)| (s, o))
+                .collect::<Vec<_>>(),
+            vec![
+                ("▁Hey", (0, 6)),
+                ("▁my", (6, 11)),
+                ("▁friend", (11, 20)),
+                ("▁", (20, 23)),
+                ("▁<s>", (23, 29)),
+                ("▁how", (29, 35)),
+                ("▁are", (35, 41)),
+                ("▁you", (41, 47))
+            ]
+        );
+    }
     #[test]
     fn decode() {
         let decoder = Metaspace::new('▁', true);

tokenizers/src/pre_tokenizers/mod.rs

@@ -104,6 +104,28 @@ mod tests {
                 PreTokenizerWrapper::Metaspace(Metaspace::new('▁', true))
             ]))
         );
+
+        let pre_tokenizer: PreTokenizerWrapper = serde_json::from_str(
+            r#"{"type":"Metaspace","replacement":"▁","add_prefix_space":true, "legacy":false}"#,
+        )
+        .unwrap();
+
+        let mut expected_pre_tokenizer = Metaspace::new('▁', true);
+        expected_pre_tokenizer.legacy = false;
+        assert_eq!(
+            pre_tokenizer,
+            PreTokenizerWrapper::Metaspace(expected_pre_tokenizer)
+        );
+
+        let pre_tokenizer: PreTokenizerWrapper = serde_json::from_str(
+            r#"{"type":"Metaspace","replacement":"▁","add_prefix_space":true, "legacy":true}"#,
+        )
+        .unwrap();
+
+        assert_eq!(
+            pre_tokenizer,
+            PreTokenizerWrapper::Metaspace(Metaspace::new('▁', true))
+        );
     }
 
     #[test]

tokenizers/src/tokenizer/mod.rs

@@ -658,6 +658,10 @@ where
 
         final_vocab
     }
+    /// Get the added tokens encoder
+    pub fn get_added_tokens_encoder(&self) -> HashMap<String, u32> {
+        self.added_vocabulary.get_vocab().clone()
+    }
 
     /// Get the added tokens decoder
     pub fn get_added_tokens_decoder(&self) -> HashMap<u32, AddedToken> {