feat: handle Merge* variants of the SplitDelimiterBehavior

tokenizers/tk-encode/src/pre_tokenizers/punctuation.rs

@@ -1,5 +1,6 @@
 use serde::{Deserialize, Serialize};
 
+use crate::pipeline;
 use crate::tokenizer::{PreTokenizedString, PreTokenizer, Result, SplitDelimiterBehavior};
 use crate::utils::macro_rules_attribute;
 use unicode_categories::UnicodeCategories;
@@ -37,6 +38,14 @@ impl PreTokenizer for Punctuation {
     }
 }
 
+impl pipeline::PreTokenizer for Punctuation {
+    fn pre_tokenize(&self, text: &str, out: &mut Vec<pipeline::Split>) -> Result<()> {
+        pipeline::split_delimiter(text, out, is_punc, self.behavior);
+
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -64,6 +73,109 @@ mod tests {
         );
     }
 
+    fn pretokenize(behavior: SplitDelimiterBehavior, text: &str) -> Vec<(&str, (u32, u32))> {
+        let pretok = Punctuation::new(behavior);
+        let mut splits = Vec::new();
+        crate::pipeline::PreTokenizer::pre_tokenize(&pretok, text, &mut splits).unwrap();
+        splits
+            .iter()
+            .map(|s| (&text[s.range()], (s.start, s.end)))
+            .collect()
+    }
+
+    #[test]
+    fn pipeline_behaviors() {
+        use SplitDelimiterBehavior::*;
+        // '-' is ASCII punctuation; these mirror the `SplitDelimiterBehavior` docs.
+        let text = "the-final--countdown";
+        assert_eq!(
+            pretokenize(Isolated, text),
+            vec![
+                ("the", (0, 3)),
+                ("-", (3, 4)),
+                ("final", (4, 9)),
+                ("-", (9, 10)),
+                ("-", (10, 11)),
+                ("countdown", (11, 20)),
+            ],
+        );
+        assert_eq!(
+            pretokenize(Removed, text),
+            vec![("the", (0, 3)), ("final", (4, 9)), ("countdown", (11, 20))],
+        );
+        assert_eq!(
+            pretokenize(Contiguous, text),
+            vec![
+                ("the", (0, 3)),
+                ("-", (3, 4)),
+                ("final", (4, 9)),
+                ("--", (9, 11)),
+                ("countdown", (11, 20)),
+            ],
+        );
+        assert_eq!(
+            pretokenize(MergedWithPrevious, text),
+            vec![
+                ("the-", (0, 4)),
+                ("final-", (4, 10)),
+                ("-", (10, 11)),
+                ("countdown", (11, 20)),
+            ],
+        );
+        assert_eq!(
+            pretokenize(MergedWithNext, text),
+            vec![
+                ("the", (0, 3)),
+                ("-final", (3, 9)),
+                ("-", (9, 10)),
+                ("-countdown", (10, 20)),
+            ],
+        );
+    }
+
+    #[test]
+    fn pipeline_matches_legacy_default() {
+        // default (Isolated) must agree with the legacy `punctuation_basic` expectation
+        assert_eq!(
+            pretokenize(
+                SplitDelimiterBehavior::Isolated,
+                "Hey friend!     How are you?!?"
+            ),
+            vec![
+                ("Hey friend", (0, 10)),
+                ("!", (10, 11)),
+                ("     How are you", (11, 27)),
+                ("?", (27, 28)),
+                ("!", (28, 29)),
+                ("?", (29, 30)),
+            ],
+        );
+    }
+
+    #[test]
+    fn pipeline_edge_cases() {
+        use SplitDelimiterBehavior::*;
+        let empty = Vec::<(&str, (u32, u32))>::new();
+        assert_eq!(pretokenize(Isolated, ""), empty);
+        assert_eq!(pretokenize(Isolated, "!"), vec![("!", (0, 1))]);
+        assert_eq!(pretokenize(Removed, "!"), empty);
+        // leading delimiter has no previous run to merge into
+        assert_eq!(
+            pretokenize(MergedWithPrevious, "-a"),
+            vec![("-", (0, 1)), ("a", (1, 2))],
+        );
+        // trailing delimiter has no following run to merge into
+        assert_eq!(
+            pretokenize(MergedWithNext, "a-"),
+            vec![("a", (0, 1)), ("-", (1, 2))],
+        );
+        // multibyte: é is 2 bytes, offsets are byte offsets
+        assert_eq!(
+            pretokenize(Isolated, "café!"),
+            vec![("café", (0, 5)), ("!", (5, 6))],
+        );
+    }
+
     #[test]
     fn deserialization() {
         let punctuation: Punctuation = serde_json::from_str(r#"{"type": "Punctuation"}"#).unwrap();

tokenizers/tk-encode/src/tokenizer/pipeline.rs

@@ -5,12 +5,12 @@ use crate::added_vocabulary::bucket_added_vocabulary::{
     AddedToken as BucketAddedToken, AddedVocabulary as BucketAddedVocabulary,
 };
 use crate::{
-    pre_tokenizers::{bert::BertPreTokenizer, whitespace::Whitespace},
+    pre_tokenizers::{bert::BertPreTokenizer, punctuation::Punctuation, whitespace::Whitespace},
     Model, ModelWrapper, NormalizedString, Normalizer, NormalizerWrapper, PostProcessorWrapper,
     PreTokenizerWrapper, Token, Tokenizer,
 };
 
-use super::Result;
+use super::{Result, SplitDelimiterBehavior};
 
 /// A pre-token split, a range into the input text.
 #[derive(Copy, Clone)]
@@ -36,6 +36,7 @@ pub trait PreTokenizer {
 /// The pre-tokenizers a [`PipelineTokenizer`] can run.
 pub enum PipelinePreTokenizer {
     Bert(BertPreTokenizer),
+    Punctuation(Punctuation),
     Whitespace(Whitespace),
     None,
 }
@@ -45,6 +46,7 @@ impl PreTokenizer for PipelinePreTokenizer {
         match self {
             Self::None => Ok(()),
             Self::Bert(pretok) => pretok.pre_tokenize(text, out),
+            Self::Punctuation(pretok) => pretok.pre_tokenize(text, out),
             Self::Whitespace(pretok) => pretok.pre_tokenize(text, out),
         }
     }
@@ -191,6 +193,7 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer {
         let pre_tokenizer = match tok.get_pre_tokenizer() {
             None => PipelinePreTokenizer::None,
             Some(PreTokenizerWrapper::BertPreTokenizer(p)) => PipelinePreTokenizer::Bert(*p),
+            Some(PreTokenizerWrapper::Punctuation(p)) => PipelinePreTokenizer::Punctuation(*p),
             Some(PreTokenizerWrapper::Whitespace(p)) => PipelinePreTokenizer::Whitespace(p.clone()),
             Some(other) => {
                 return Err(format!(
@@ -386,3 +389,73 @@ pub fn split<C: Copy + PartialEq>(
         }
     }
 }
+
+/// Splits `text` around a single delimiter predicate, honoring the full
+/// [`SplitDelimiterBehavior`] contract. The pipeline-side equivalent of
+/// `NormalizedString::split(pattern, behavior)` for a char predicate.
+///
+/// The three non-merging behaviors reduce to a [`SplitPolicy`] on the delimiter
+/// class and reuse [`split`]. The two merge variants are their own single pass:
+/// - `MergedWithPrevious` cuts the split *after* each delimiter, so a delimiter
+///   joins the run before it (`"the-final"` -> `["the-", "final"]`).
+/// - `MergedWithNext` cuts *before* each delimiter, so it joins the run after it
+///   (`"the-final"` -> `["the", "-final"]`).
+pub fn split_delimiter(
+    text: &str,
+    out: &mut Vec<Split>,
+    is_delim: impl Fn(char) -> bool,
+    behavior: SplitDelimiterBehavior,
+) {
+    use SplitDelimiterBehavior::*;
+
+    let delim_policy = match behavior {
+        Removed => SplitPolicy::Remove,
+        Isolated => SplitPolicy::Isolate,
+        Contiguous => SplitPolicy::Keep,
+        MergedWithPrevious => {
+            let mut start: u32 = 0;
+            for (i, ch) in text.char_indices() {
+                if is_delim(ch) {
+                    let end = (i + ch.len_utf8()) as u32;
+                    out.push(Split { start, end });
+                    start = end;
+                }
+            }
+            if (start as usize) < text.len() {
+                out.push(Split {
+                    start,
+                    end: text.len() as u32,
+                });
+            }
+            return;
+        }
+        MergedWithNext => {
+            let mut start: u32 = 0;
+            for (i, ch) in text.char_indices() {
+                if is_delim(ch) {
+                    let i = i as u32;
+                    // skip the empty span before a leading run of delimiters
+                    if i > start {
+                        out.push(Split { start, end: i });
+                    }
+                    start = i;
+                }
+            }
+            if (start as usize) < text.len() {
+                out.push(Split {
+                    start,
+                    end: text.len() as u32,
+                });
+            }
+            return;
+        }
+    };
+
+    split(text, out, is_delim, |d| {
+        if d {
+            delim_policy
+        } else {
+            SplitPolicy::Keep
+        }
+    });
+}