Fix Phi 4 tokenization (#153)

* Fix Phi 4 tokenization

* Fix: remove empty trailing pretokenized match

* Realistic pre-tokenizer tests, double-checked with transformers.js

* Couple more test cases

Co-authored-by: Anthony <anthony@depasquale.org>

Author: pcuenca

Sources/Tokenizers/PreTokenizer.swift

@@ -248,7 +248,7 @@ extension StringSplitPattern {
     func split(_ text: String, invert: Bool = true) -> [String] {
         switch self {
         case .regexp(let regexp):
-            return text.split(by: regexp, includeSeparators: !invert)
+            return text.split(by: regexp, includeSeparators: true)
         case .string(let substring):
             return text.split(by: substring, options: [], includeSeparators: !invert)
         }
@@ -292,7 +292,9 @@ public extension String {
             start = range.upperBound
         }
 
-        result.append(String(self[start...]))
+        if omittingEmptySubsequences && start < endIndex {
+            result.append(String(self[start...]))
+        }
         return result
     }
 

Tests/PreTokenizerTests/PreTokenizerTests.swift

@@ -119,7 +119,7 @@ class PreTokenizerTests: XCTestCase {
         )
         XCTAssertEqual(
             preTokenizer1.preTokenize(text: "   Hey,    friend,    what's up?  "),
-            [" ", " ", " ", "Hey,", " ", " ", " ", " ", "friend,", " ", " ", " ", " ", "what's", " ", "up?", " ", " ", ""]
+            [" ", " ", " ", "Hey,", " ", " ", " ", " ", "friend,", " ", " ", " ", " ", "what's", " ", "up?", " ", " "]
         )
 
         let preTokenizer2 = SplitPreTokenizer(config: Config(["pattern": ["Regex": "\\s"]]))
@@ -133,21 +133,22 @@ class PreTokenizerTests: XCTestCase {
         )
         XCTAssertEqual(
             preTokenizer2.preTokenize(text: "   Hey,    friend,    what's up?  "),
-            [" ", " ", " ", "Hey,", " ", " ", " ", " ", "friend,", " ", " ", " ", " ", "what's", " ", "up?", " ", " ", ""]
+            [" ", " ", " ", "Hey,", " ", " ", " ", " ", "friend,", " ", " ", " ", " ", "what's", " ", "up?", " ", " "]
         )
 
-        let preTokenizer3 = SplitPreTokenizer(config: Config(["pattern": ["Regex": "\\s"], "invert": true]))
+        let preTokenizer3 = SplitPreTokenizer(config: Config(["pattern": ["Regex": "(?i:\'s|\'t|\'re|\'ve|\'m|\'ll|\'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"], "invert": true]))
         XCTAssertEqual(
-            preTokenizer3.preTokenize(text: "Hey friend!"),
-            ["Hey", "friend!"]
+            preTokenizer3.preTokenize(text: "Hello"),
+            ["Hello"]
         )
+
         XCTAssertEqual(
-            preTokenizer3.preTokenize(text: "Hey friend!     How are you?!?"),
-            ["Hey", "friend!", "How", "are", "you?!?"]
+            preTokenizer3.preTokenize(text: "Hey friend!"),
+            ["Hey", " friend", "!"]
         )
         XCTAssertEqual(
-            preTokenizer3.preTokenize(text: "   Hey,    friend,    what's up?  "),
-            ["Hey,", "friend,", "what's", "up?", ""]
+            preTokenizer3.preTokenize(text: "Hey friend!     How are you?!?"),
+            ["Hey", " friend", "!", "    ", " How", " are", " you", "?!?"]
         )
     }
 

Tests/TokenizersTests/TokenizerTests.swift

@@ -95,6 +95,19 @@ class GemmaUnicodeTests: XCTestCase {
     }
 }
 
+class PhiSimpleTests: XCTestCase {
+    func testPhi4() async throws {
+        guard let tokenizer = try await AutoTokenizer.from(pretrained: "microsoft/phi-4") as? PreTrainedTokenizer else {
+            XCTFail()
+            return
+        }
+
+        XCTAssertEqual(tokenizer.encode(text: "hello"), [15339])
+        XCTAssertEqual(tokenizer.encode(text: "hello world"), [15339, 1917])
+        XCTAssertEqual(tokenizer.encode(text: "<|im_start|>user<|im_sep|>Who are you?<|im_end|><|im_start|>assistant<|im_sep|>"), [100264, 882, 100266, 15546, 527, 499, 30, 100265, 100264, 78191, 100266])
+    }
+}
+
 
 struct EncodedTokenizerSamplesDataset: Decodable {
     let text: String