Add conditional normalization

Systemcluster · Dec 20, 2024 · f2e7d32 · f2e7d32
1 parent 7ceb134
commit f2e7d32
Show file tree

Hide file tree

Showing 57 changed files with 1,420,337 additions and 117 deletions.
diff --git a/src/config.rs b/src/config.rs
@@ -2,6 +2,8 @@
 //!
 //! Defines the tokenization mode fallback, input normalization, pre-tokenization split behavior, post-tokenization processing, post-decode processing, and input templates.
 
+use core::ops::Range;
+
 use alloc::borrow::Cow;
 use alloc::string::String;
 use alloc::vec::Vec;
@@ -128,12 +130,12 @@ impl Configuration {
 
     /// Normalizes the input before tokenization.
     #[inline(never)]
-    pub fn normalize(&self, text: &mut Cow<str>) {
+    pub fn normalize(&self, text: &mut Cow<str>, position: Range<usize>) {
         if text.is_empty() {
             return;
         }
         for norm in &self.normalization {
-            norm.normalize(text);
+            norm.normalize(text, position.clone());
         }
     }
 

diff --git a/src/config/normalization.rs b/src/config/normalization.rs
@@ -1,5 +1,7 @@
 //! Pre-tokenization input normalization.
 
+use core::ops::Range;
+
 use alloc::borrow::Cow;
 use alloc::boxed::Box;
 use alloc::string::{String, ToString};
@@ -62,6 +64,14 @@ impl From<Regex> for NormalizationReplacePattern {
     }
 }
 
+/// Condition for conditional normalization.
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "serialization", derive(Deserialize, Serialize))]
+pub enum NormalizationCondition {
+    StartOfText,
+    EndOfText,
+}
+
 /// Pre-tokenization input normalization configuration.
 #[derive(Debug, Clone, PartialEq, Eq)]
 #[cfg_attr(feature = "serialization", derive(Deserialize, Serialize))]
@@ -98,11 +108,16 @@ pub enum Normalization {
     },
     /// Precompiled character map.
     CharsMap { map: CharsMap },
+    /// Conditional normalization.
+    Conditional {
+        condition:     NormalizationCondition,
+        normalization: Box<Normalization>,
+    },
 }
 
 impl Normalization {
     #[inline(never)]
-    pub fn normalize(&self, text: &mut Cow<str>) {
+    pub fn normalize(&self, text: &mut Cow<str>, position: Range<usize>) {
         use Normalization::*;
         match self {
             Unicode { scheme } => {
@@ -147,6 +162,17 @@ impl Normalization {
             CharsMap { map } => {
                 normalize_charsmap(text, map);
             }
+            Conditional {
+                condition,
+                normalization,
+            } => {
+                if match condition {
+                    NormalizationCondition::StartOfText => position.start == 0,
+                    NormalizationCondition::EndOfText => position.end == usize::MAX,
+                } {
+                    normalization.normalize(text, position);
+                }
+            }
         }
     }
 }
@@ -324,20 +350,20 @@ mod tests {
     fn test_normalization_nmt() {
         let mut text = Cow::Borrowed("aaa\u{200D}bbb\u{8f}");
         let normalization = Normalization::NMT;
-        normalization.normalize(&mut text);
+        normalization.normalize(&mut text, 0..usize::MAX);
         assert_eq!(text, "aaa bbb");
     }
 
     #[test]
     fn test_normalization_case_fold() {
         let mut text = Cow::Borrowed("AAA bbb");
         let normalization = Normalization::CaseFold { upper: false };
-        normalization.normalize(&mut text);
+        normalization.normalize(&mut text, 0..usize::MAX);
         assert_eq!(text, "aaa bbb");
 
         let mut text = Cow::Borrowed("AAA bbb");
         let normalization = Normalization::CaseFold { upper: true };
-        normalization.normalize(&mut text);
+        normalization.normalize(&mut text, 0..usize::MAX);
         assert_eq!(text, "AAA BBB");
     }
 
@@ -347,7 +373,7 @@ mod tests {
         let normalization = Normalization::Append {
             append: " bbb".to_string(),
         };
-        normalization.normalize(&mut text);
+        normalization.normalize(&mut text, 0..usize::MAX);
         assert_eq!(text, "aaa bbb");
     }
 
@@ -357,7 +383,7 @@ mod tests {
         let normalization = Normalization::Prepend {
             prepend: "aaa ".to_string(),
         };
-        normalization.normalize(&mut text);
+        normalization.normalize(&mut text, 0..usize::MAX);
         assert_eq!(text, "aaa bbb");
     }
 
@@ -370,7 +396,7 @@ mod tests {
             right:     3,
             pad:       false,
         };
-        normalization.normalize(&mut text);
+        normalization.normalize(&mut text, 0..usize::MAX);
         assert_eq!(text, "aabbbaaa");
 
         let mut text = Cow::Borrowed("aba");
@@ -380,7 +406,7 @@ mod tests {
             right:     3,
             pad:       true,
         };
-        normalization.normalize(&mut text);
+        normalization.normalize(&mut text, 0..usize::MAX);
         assert_eq!(text, "aabaaa");
     }
 
@@ -392,15 +418,15 @@ mod tests {
             left:      2,
             right:     3,
         };
-        normalization.normalize(&mut text);
+        normalization.normalize(&mut text, 0..usize::MAX);
         assert_eq!(text, "aba");
     }
 
     #[test]
     fn test_normalization_collapse() {
         let mut text = Cow::Borrowed("abbbba bbb");
         let normalization = Normalization::Collapse { character: 'b' };
-        normalization.normalize(&mut text);
+        normalization.normalize(&mut text, 0..usize::MAX);
         assert_eq!(text, "aba b");
     }
 
@@ -411,7 +437,40 @@ mod tests {
             pattern:     Regex::new(r"b").unwrap().into(),
             replacement: "a".to_string(),
         };
-        normalization.normalize(&mut text);
+        normalization.normalize(&mut text, 0..usize::MAX);
         assert_eq!(text, "aaa aaa");
     }
+
+    #[test]
+    fn test_normalization_conditional() {
+        let mut text = Cow::Borrowed("aba bbb");
+        let normalization = Normalization::Conditional {
+            condition:     NormalizationCondition::StartOfText,
+            normalization: Box::new(Normalization::Replace {
+                pattern:     Regex::new(r"b").unwrap().into(),
+                replacement: "a".to_string(),
+            }),
+        };
+        normalization.normalize(&mut text, 0..usize::MAX);
+        assert_eq!(text, "aaa aaa");
+
+        let mut text = Cow::Borrowed("aba bbb");
+        normalization.normalize(&mut text, 1..usize::MAX);
+        assert_eq!(text, "aba bbb");
+
+        let mut text = Cow::Borrowed("aba bbb");
+        let normalization = Normalization::Conditional {
+            condition:     NormalizationCondition::EndOfText,
+            normalization: Box::new(Normalization::Replace {
+                pattern:     Regex::new(r"b").unwrap().into(),
+                replacement: "a".to_string(),
+            }),
+        };
+        normalization.normalize(&mut text, 0..usize::MAX);
+        assert_eq!(text, "aaa aaa");
+
+        let mut text = Cow::Borrowed("aba bbb");
+        normalization.normalize(&mut text, 0..4);
+        assert_eq!(text, "aba bbb");
+    }
 }
diff --git a/src/convert/tekken.rs b/src/convert/tekken.rs
@@ -75,7 +75,7 @@ use ms::Tokenizer;
 /// use kitoken::convert::convert_tekken;
 /// use kitoken::Kitoken;
 ///
-/// let data = std::fs::read("tests/models/tekken/nemo.json")?;
+/// let data = std::fs::read("tests/models/tekken/mistral2410.json")?;
 /// let definition = convert_tekken(data).unwrap();
 ///
 /// let tokenizer = Kitoken::try_from(definition).unwrap();

diff --git a/src/convert/tokenizers.rs b/src/convert/tokenizers.rs
@@ -17,9 +17,9 @@ use hashbrown::HashMap;
 use crate::convert::ConversionError;
 use crate::{
     Configuration, Decoding, Definition, Fallback, InsertionPosition, Kitoken, Metadata, Model,
-    Normalization, Processing, ProcessingDirection, Regex, Scores, SpecialToken, SpecialTokenKind,
-    SpecialVocab, Split, SplitBehavior, Template, Token, TokenBytes, TokenId, UnicodeNormalization,
-    Vocab,
+    Normalization, NormalizationCondition, Processing, ProcessingDirection, Regex, Scores,
+    SpecialToken, SpecialTokenKind, SpecialVocab, Split, SplitBehavior, Template, Token,
+    TokenBytes, TokenId, UnicodeNormalization, Vocab,
 };
 
 mod hf {
@@ -642,13 +642,23 @@ pub fn convert_tokenizers(data: impl AsRef<[u8]>) -> Result<Definition, Conversi
                     pattern:     Regex::new(r" ")?.into(),
                     replacement: replacement.to_string(),
                 });
-                if prepend_scheme != PrependScheme::Never {
+                if prepend_scheme == PrependScheme::Always {
                     config.normalization.push(Normalization::Extend {
                         character: replacement,
                         left:      1,
                         right:     0,
                         pad:       true,
                     });
+                } else if prepend_scheme == PrependScheme::First {
+                    config.normalization.push(Normalization::Conditional {
+                        condition:     NormalizationCondition::StartOfText,
+                        normalization: Box::new(Normalization::Extend {
+                            character: replacement,
+                            left:      1,
+                            right:     0,
+                            pad:       true,
+                        }),
+                    });
                 }
                 if split {
                     config.split.push(Split::Pattern {

diff --git a/src/lib.rs b/src/lib.rs
@@ -277,7 +277,7 @@ impl Kitoken {
             if let Some(next) = extracted.pop() {
                 if next.0 > posit {
                     let mut text = text[posit..next.0].into();
-                    self.config.normalize(&mut text);
+                    self.config.normalize(&mut text, posit..next.0);
                     parts.push(TextPart {
                         text,
                         special: Token::INVALID,
@@ -295,7 +295,7 @@ impl Kitoken {
                 posit = next.1;
             } else {
                 let mut rest = text[posit..text.len()].into();
-                self.config.normalize(&mut rest);
+                self.config.normalize(&mut rest, posit..usize::MAX);
                 parts.push(TextPart {
                     text:    rest,
                     special: Token::INVALID,