fix(Tokenizers): avoid quadratic regex matching for added tokens without whitespace padding

Sources/Tokenizers/Tokenizer.swift

@@ -526,9 +526,10 @@ public class PreTrainedTokenizer: @unchecked Sendable, Tokenizer {
             let token = NSRegularExpression.escapedPattern(for: $0.content)
             let prefix = $0.prefix ? #"\s*"# : ""
             let suffix = $0.suffix ? #"\s*"# : ""
+            guard $0.prefix || $0.suffix else { return "(?:\(token))" }
             return "\(prefix)(\(token))\(suffix)"
         }.joined(separator: "|")
-        addedTokensRegex = try? NSRegularExpression(pattern: addedTokensRegexString, options: [])
+        addedTokensRegex = unwrappedAddedTokens.isEmpty ? nil : try? NSRegularExpression(pattern: addedTokensRegexString, options: [])
 
         self.specialTokens = specialTokens
         self.addedTokens = Set(addedTokens.keys)

Tests/TokenizersTests/TokenizerTests.swift

@@ -38,4 +38,269 @@ struct TokenizerTests {
         let encoded = tokenizer.encode(text: "offline path")
         #expect(!encoded.isEmpty)
     }
+
+    /// https://github.com/huggingface/swift-transformers/issues/96
+    @Test
+    func legacyLlamaBehaviour() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "mlx-community/Phi-3-mini-4k-instruct-4bit-no-q-embed") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        let inputIds = tokenizer(" Hi")
+        #expect(inputIds == [1, 29871, 6324])
+    }
+
+    /// https://github.com/huggingface/swift-transformers/issues/352 (Bug 3)
+    @Test
+    func t5UnigramKeycapEmoji() async throws {
+        // Keycap "1️⃣" is one grapheme cluster but three Unicode scalars:
+        //   U+0031 DIGIT ONE
+        //   U+FE0F VARIATION SELECTOR-16
+        //   U+20E3 COMBINING ENCLOSING KEYCAP
+        // SentencePiece Unigram lookups operate per scalar, so HF Python tokenizes
+        // this as ▁1 <unk> </s> — the digit `1` matches; the VS-16 + keycap tail UNKs.
+        // Pre-fix, the Swift tokenizer iterated by `Character`, so the whole grapheme
+        // became one lattice slot that never matched any vocab entry and the digit was
+        // silently lost (▁ <unk> </s>).
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "google-t5/t5-small") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        #expect(tokenizer.encode(text: "1️⃣") == [209, 2, 1])
+    }
+    /// https://github.com/huggingface/swift-transformers/issues/352 (Bug 4)
+    @Test
+    func llama7bCombiningMarks() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "huggyllama/llama-7b") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        #expect(tokenizer.encode(text: "สวัส") == [1, 29871, 30547, 30492, 30510, 30547])
+    }
+
+    /// https://github.com/huggingface/swift-transformers/issues/99
+    @Test
+    func robertaXLMTokenizer() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "intfloat/multilingual-e5-small") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        let ids = tokenizer.encode(text: "query: how much protein should a female eat")
+        let expected = [0, 41, 1294, 12, 3642, 5045, 21308, 5608, 10, 117776, 73203, 2]
+        #expect(ids == expected)
+    }
+
+    /// https://github.com/huggingface/swift-transformers/issues/318
+    @Test
+    func kredorPunctuateAllTokenizer() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "kredor/punctuate-all") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        let ids = tokenizer.encode(text: "okay so lets get started")
+        let expected = [0, 68403, 221, 2633, 7, 2046, 26859, 2]
+        #expect(ids == expected)
+    }
+
+    @Test
+    func robertaXLMCanonicalTokenizer() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "FacebookAI/xlm-roberta-base") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        let ids = tokenizer.encode(text: "okay so lets get started")
+        let expected = [0, 68403, 221, 2633, 7, 2046, 26859, 2]
+        #expect(ids == expected)
+    }
+
+    @Test
+    func nllbTokenizer() async throws {
+        do {
+            _ = try await AutoTokenizer.from(pretrained: "Xenova/nllb-200-distilled-600M")
+            Issue.record("Expected AutoTokenizer.from to throw for strict mode")
+        } catch {
+            // Expected to throw in normal (strict) mode
+        }
+
+        // no strict mode proceeds
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "Xenova/nllb-200-distilled-600M", strict: false) as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        let ids = tokenizer.encode(text: "Why did the chicken cross the road?")
+        let expected = [256047, 24185, 4077, 349, 1001, 22690, 83580, 349, 82801, 248130, 2]
+        #expect(ids == expected)
+    }
+
+    /// Deepseek needs a post-processor override to add a bos token as in the reference implementation
+    @Test
+    func deepSeekPostProcessor() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+        #expect(tokenizer.encode(text: "Who are you?") == [151646, 15191, 525, 498, 30])
+    }
+
+    /// Some Llama tokenizers already use a bos-prepending Template post-processor
+    @Test
+    func llamaPostProcessor() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "coreml-projects/Llama-2-7b-chat-coreml") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+        #expect(tokenizer.encode(text: "Who are you?") == [1, 11644, 526, 366, 29973])
+    }
+
+    @Test
+    func localTokenizerFromPretrained() async throws {
+        let downloadDestination: URL = {
+            let base = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first!
+            return base.appending(component: "hf-local-pretrained-tests-downloads")
+        }()
+        if FileManager.default.fileExists(atPath: downloadDestination.path) {
+            try FileManager.default.removeItem(at: downloadDestination)
+        }
+        defer {
+            try? FileManager.default.removeItem(at: downloadDestination)
+        }
+
+        let hubApi = HubApi(downloadBase: downloadDestination)
+        let downloadedTo = try await hubApi.snapshot(from: "pcuenq/gemma-tokenizer")
+
+        let tokenizerOpt = try await AutoTokenizer.from(modelFolder: downloadedTo) as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+    }
+
+    @Test
+    func bertCased() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "distilbert/distilbert-base-multilingual-cased") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        #expect(tokenizer.encode(text: "mąka") == [101, 181, 102075, 10113, 102])
+        #expect(tokenizer.tokenize(text: "Car") == ["Car"])
+    }
+
+    @Test
+    func bertCasedResaved() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "pcuenq/distilbert-base-multilingual-cased-tokenizer") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        #expect(tokenizer.encode(text: "mąka") == [101, 181, 102075, 10113, 102])
+    }
+
+    @Test
+    func bertUncased() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "google-bert/bert-base-uncased") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        #expect(tokenizer.tokenize(text: "mąka") == ["ma", "##ka"])
+        #expect(tokenizer.encode(text: "mąka") == [101, 5003, 2912, 102])
+        #expect(tokenizer.tokenize(text: "département") == ["depart", "##ement"])
+        #expect(tokenizer.encode(text: "département") == [101, 18280, 13665, 102])
+        #expect(tokenizer.tokenize(text: "Car") == ["car"])
+
+        #expect(tokenizer.tokenize(text: "€4") == ["€", "##4"])
+        #expect(tokenizer.tokenize(text: "test $1 R2 #3 €4 £5 ¥6 ₣7 ₹8 ₱9 test") == ["test", "$", "1", "r", "##2", "#", "3", "€", "##4", "£5", "¥", "##6", "[UNK]", "₹", "##8", "₱", "##9", "test"])
+
+        let text = "l'eure"
+        let tokenized = tokenizer.tokenize(text: text)
+        #expect(tokenized == ["l", "'", "eu", "##re"])
+        let encoded = tokenizer.encode(text: text)
+        #expect(encoded == [101, 1048, 1005, 7327, 2890, 102])
+        let decoded = tokenizer.decode(tokens: encoded, skipSpecialTokens: true)
+        // Note: this matches the behaviour of the Python "slow" tokenizer, but the fast one produces "l ' eure"
+        #expect(decoded == "l'eure")
+
+        // Reading added_tokens from tokenizer.json
+        #expect(tokenizer.convertTokenToId("[PAD]") == 0)
+        #expect(tokenizer.convertTokenToId("[UNK]") == 100)
+        #expect(tokenizer.convertTokenToId("[CLS]") == 101)
+        #expect(tokenizer.convertTokenToId("[SEP]") == 102)
+        #expect(tokenizer.convertTokenToId("[MASK]") == 103)
+    }
+
+    @Test
+    func bertHangul() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "BAAI/bge-small-en-v1.5") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        #expect(tokenizer.encode(text: "안") == [101, 1463, 30006, 30021, 102])
+    }
+
+    @Test
+    func bertJapaneseDakuten() async throws {
+        // Regression for https://github.com/huggingface/swift-transformers/issues/352 (Bug 2).
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "BAAI/bge-small-en-v1.5") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        // Minimal repro: a single voiced-kana character.
+        #expect(tokenizer.encode(text: "ザ") == [101, 1705, 102])
+
+        // Long form covering hiragana, ideographs, punctuation, and katakana with
+        // multiple voiced kana (`ザ` and `で`).
+        #expect(
+            tokenizer.encode(text: "こんにちは、世界。トークナイザーのテストです。") == [
+                101, 1655, 30217, 30194, 30188, 30198, 1635, 1745, 100, 1636,
+                1714, 30265, 30228, 30241, 30221, 30231, 30265, 30197, 30239,
+                30233, 30240, 30191, 30184, 1636, 102,
+            ])
+    }
+
+    @Test
+    func robertaEncodeDecode() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "FacebookAI/roberta-base") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        #expect(tokenizer.tokenize(text: "l'eure") == ["l", "'", "e", "ure"])
+        #expect(tokenizer.encode(text: "l'eure") == [0, 462, 108, 242, 2407, 2])
+        #expect(tokenizer.decode(tokens: tokenizer.encode(text: "l'eure"), skipSpecialTokens: true) == "l'eure")
+
+        #expect(tokenizer.tokenize(text: "mąka") == ["m", "Ä", "ħ", "ka"])
+        #expect(tokenizer.encode(text: "mąka") == [0, 119, 649, 5782, 2348, 2])
+
+        #expect(tokenizer.tokenize(text: "département") == ["d", "é", "part", "ement"])
+        #expect(tokenizer.encode(text: "département") == [0, 417, 1140, 7755, 6285, 2])
+
+        #expect(tokenizer.tokenize(text: "Who are you?") == ["Who", "Ġare", "Ġyou", "?"])
+        #expect(tokenizer.encode(text: "Who are you?") == [0, 12375, 32, 47, 116, 2])
+
+        #expect(tokenizer.tokenize(text: " Who are you? ") == ["ĠWho", "Ġare", "Ġyou", "?", "Ġ"])
+        #expect(tokenizer.encode(text: " Who are you? ") == [0, 3394, 32, 47, 116, 1437, 2])
+
+        #expect(tokenizer.tokenize(text: "<s>Who are you?</s>") == ["<s>", "Who", "Ġare", "Ġyou", "?", "</s>"])
+        #expect(tokenizer.encode(text: "<s>Who are you?</s>") == [0, 0, 12375, 32, 47, 116, 2, 2])
+    }
+
+    @Test
+    func tokenizerBackend() async throws {
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "mlx-community/Ministral-3-3B-Instruct-2512-4bit") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        #expect(tokenizer.encode(text: "She took a train to the West") == [6284, 5244, 1261, 10018, 1317, 1278, 5046])
+    }
+
+    @Test
+    func addedTokensRegexPerformanceAndCorrectness() async throws {
+        // Regression coverage for #383: Tokenizer.encode was quadratic in the number of added tokens
+        // due to capturing groups in NSRegularExpression. Non-capturing groups must preserve exact
+        // tokenization for standard added tokens while allowing prefix/suffix stripped tokens to capture correctly.
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: "pcuenq/gemma-tokenizer") as? PreTrainedTokenizer
+        #expect(tokenizerOpt != nil)
+        let tokenizer = tokenizerOpt!
+
+        // Verify that added tokens and newlines are matched and encoded properly without quadratic explosion
+        let multiLineText = String(repeating: "Hello world\n", count: 20)
+        let encoded = tokenizer.encode(text: multiLineText)
+        #expect(!encoded.isEmpty)
+
+        // Verify roundtrip decoding
+        let decoded = tokenizer.decode(tokens: encoded, skipSpecialTokens: false)
+        #expect(decoded.contains("Hello world"))
+    }
 }