Split by regexp with capture groups

The other split helpers we have don't work for capture groups.
We had to resort to raw `NSRegularExpression`s

Author: pcuenca

Sources/Tokenizers/PreTokenizer.swift

@@ -248,7 +248,7 @@ extension StringSplitPattern {
     }
 }
 
-extension String {
+public extension String {
     func ranges(of string: String, options: CompareOptions = .regularExpression) -> [Range<Index>] {
         var result: [Range<Index>] = []
         var start = startIndex
@@ -277,6 +277,42 @@ extension String {
         return result
     }
 
+    /// This version supports capture groups, wheres the one above doesn't
+    func split(by captureRegex: NSRegularExpression) -> [String] {
+        // Find the matching capture groups
+        let selfRange = NSRange(startIndex..<endIndex, in: self)
+        let matches = captureRegex.matches(in: self, options: [], range: selfRange)
+
+        if matches.first == nil { return [self] }
+
+        var result: [String] = []
+        var start = startIndex
+        for match in matches {
+            // Append prefix before matched separator
+            let prefixEnd = index(startIndex, offsetBy: match.range.lowerBound)
+            if start < prefixEnd {
+                result.append(String(self[start..<prefixEnd]))
+            }
+            start = index(startIndex, offsetBy: match.range.upperBound)
+
+            // Append separator, supporting capture groups
+            for r in (0..<match.numberOfRanges).reversed() {
+                let matchRange = match.range(at: r)
+                if let sepRange = Range(matchRange, in:self) {
+                    result.append(String(self[sepRange]))
+                    break
+                }
+            }
+        }
+
+        // Append remaining suffix
+        let beginningOfEnd = index(startIndex, offsetBy: matches.last!.range.upperBound)
+        if beginningOfEnd < endIndex {
+            result.append(String(self[beginningOfEnd...]))
+        }
+
+        return result
+    }
 }
 
 public enum SplitDelimiterBehavior {

Tests/TokenizersTests/AddedTokensTests.swift

@@ -0,0 +1,67 @@
+//
+//  AddedTokensTests.swift
+//
+//
+//  Created by Pedro Cuenca on 20240426.
+//
+
+import XCTest
+import Tokenizers
+import Hub
+
+class AddedTokensTests: XCTestCase {
+    func testPhiAddedEnd() async throws {
+        let tokenizer = try await AutoTokenizer.from(pretrained: "mlx-community/Phi-3-mini-128k-instruct-4bit")
+        let inputIds = tokenizer("This is the <|end|>. My only friend, the <|end|>")
+        XCTAssertEqual(inputIds, [1, 910, 338, 278, 29871, 32007, 29889, 1619, 871, 5121, 29892, 278, 29871, 32007])
+
+        let decoded = tokenizer.decode(tokens: inputIds)
+        XCTAssertEqual(decoded, "<s> This is the <|end|>. My only friend, the <|end|>")
+    }
+
+    func testSplitWithCaptureGroups() {
+        let addedTokensRegexp = #"(<\|end\|>)\s*|(<\|raw\|>)\s*"#
+        let captureRegex = try! NSRegularExpression(pattern: addedTokensRegexp, options: [])
+
+        XCTAssertEqual(
+            "eating <|raw|> meat <|end|> That's all".split(by: captureRegex),
+            ["eating ", "<|raw|>", "meat ", "<|end|>", "That's all"]
+        )
+
+        XCTAssertEqual(
+            "<|raw|>".split(by: captureRegex),
+            ["<|raw|>"]
+        )
+
+        XCTAssertEqual(
+            "This string doesn't have those separators".split(by: captureRegex),
+            ["This string doesn't have those separators"]
+        )
+
+        XCTAssertEqual(
+            "start <|end|>".split(by: captureRegex),
+            ["start ", "<|end|>"]
+        )
+
+        XCTAssertEqual(
+            "start <|end|> ".split(by: captureRegex),
+            ["start ", "<|end|>"]
+        )
+
+        XCTAssertEqual(
+            "start <|end|>       ".split(by: captureRegex),
+            ["start ", "<|end|>"]
+        )
+
+        XCTAssertEqual(
+            "start <|end|>       for real".split(by: captureRegex),
+            ["start ", "<|end|>", "for real"]
+        )
+
+        XCTAssertEqual(
+            "<|raw|><|end|>".split(by: captureRegex),
+            ["<|raw|>", "<|end|>"]
+        )
+
+    }
+}