From d4f827b7e2bf6b32a910cae036ad84c85b613960 Mon Sep 17 00:00:00 2001 From: michalharakal Date: Sat, 24 Oct 2026 09:36:31 +0200 Subject: [PATCH] tokenizer: index added tokens by first char and length AddedTokenMatcher compared every added token against every input position. Gemma's GGUF declares ~6,400 added tokens, so one 70-character chat turn cost ~450k region compares. A token can only match at a position if it starts with that character, and at most one token per length can match, so the matcher now keeps the distinct lengths per first character and does one map lookup per length. Same results (longest first, single_word falls back to shorter); 2.0 ms -> 36 us per encode on the JVM with a Gemma-3 270M vocab, ids identical. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 4 ++++ .../ainet/frontend/tokenizer/TextPipeline.kt | 19 +++++++++++++------ .../tokenizer/SpecialTokenSplitterTest.kt | 15 +++++++++++++++ 3 files changed, 32 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a84b7548..fd3b473d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- **Added-token matching no longer scans every added token at every position** (`skainet-frontend-text`, `AddedTokenMatcher` behind `SpecialTokenSplitter` and the inline matchers). Gemma's GGUF declares ~6,400 added tokens; a 70-character chat turn cost ~450k region compares per `encode`. Candidates are now indexed by first character and length (one map lookup per distinct length). Gemma-3 270M fine-tune vocab, JVM: 2.0 ms → 36 µs per chat-turn encode, ids unchanged; on a 32-bit ARM Android device the same encode took 36–166 ms. + ## [0.62.1] - 2026-10-21 Headline: **the native Q8_0 / Q4_0 / Q5_0 / Q5_1 kernel bindings no longer zero the earlier rows of a diff --git a/skainet-frontend/skainet-frontend-text/src/commonMain/kotlin/sk/ainet/frontend/tokenizer/TextPipeline.kt b/skainet-frontend/skainet-frontend-text/src/commonMain/kotlin/sk/ainet/frontend/tokenizer/TextPipeline.kt index f00e19ca..47b6ea4e 100644 --- a/skainet-frontend/skainet-frontend-text/src/commonMain/kotlin/sk/ainet/frontend/tokenizer/TextPipeline.kt +++ b/skainet-frontend/skainet-frontend-text/src/commonMain/kotlin/sk/ainet/frontend/tokenizer/TextPipeline.kt @@ -37,9 +37,16 @@ public interface RunDecoder { */ internal class AddedTokenMatcher(tokens: Collection) { private val byContent: Map = tokens.filter { it.content.isNotEmpty() }.associateBy { it.content } - private val byLengthDesc: List = byContent.keys.sortedByDescending { it.length } + /** + * The distinct lengths of the added tokens per first character, longest first. Only a token starting with + * `text[from]` can match at `from`, and at most one per length, so a position costs one map lookup per + * length. Gemma's GGUF declares ~6,400 added tokens; comparing each of them at every position cost ~450k + * region compares for a 70-character chat turn. + */ + private val lengthsByFirstChar: Map = + byContent.keys.groupBy { it[0] }.mapValues { (_, c) -> c.map { it.length }.distinct().sortedDescending().toIntArray() } - val isEmpty: Boolean get() = byLengthDesc.isEmpty() + val isEmpty: Boolean get() = byContent.isEmpty() val ids: Set = byContent.values.map { it.id }.toSet() val specialIds: Set = byContent.values.filter { it.special }.map { it.id }.toSet() val idToContent: Map = byContent.values.associate { it.id to it.content } @@ -64,12 +71,12 @@ internal class AddedTokenMatcher(tokens: Collection) { } private fun matchAt(text: String, from: Int): AddedToken? { - for (content in byLengthDesc) { - if (!text.regionMatches(from, content, 0, content.length)) continue - val tok = byContent.getValue(content) + for (length in lengthsByFirstChar[text[from]] ?: return null) { + if (from + length > text.length) continue + val tok = byContent[text.substring(from, from + length)] ?: continue if (tok.singleWord) { val before = if (from > 0) text[from - 1] else ' ' - val afterIdx = from + content.length + val afterIdx = from + length val after = if (afterIdx < text.length) text[afterIdx] else ' ' if (before.isLetterOrDigit() || after.isLetterOrDigit()) continue } diff --git a/skainet-frontend/skainet-frontend-text/src/commonTest/kotlin/sk/ainet/frontend/tokenizer/SpecialTokenSplitterTest.kt b/skainet-frontend/skainet-frontend-text/src/commonTest/kotlin/sk/ainet/frontend/tokenizer/SpecialTokenSplitterTest.kt index b7f00c10..645c55c4 100644 --- a/skainet-frontend/skainet-frontend-text/src/commonTest/kotlin/sk/ainet/frontend/tokenizer/SpecialTokenSplitterTest.kt +++ b/skainet-frontend/skainet-frontend-text/src/commonTest/kotlin/sk/ainet/frontend/tokenizer/SpecialTokenSplitterTest.kt @@ -83,6 +83,21 @@ class SpecialTokenSplitterTest { assertEquals("a<|im_start|>b", splitter.decode(ids)) } + @Test + fun `a rejected single_word match falls back to the next shorter token with the same first char`() { + val splitter = SpecialTokenSplitter( + CharBase, + listOf( + AddedToken(content = "word", id = 5, special = false, singleWord = true), + AddedToken(content = "wo", id = 6, special = false), + AddedToken(content = "wordy-long", id = 9, special = false), // longer than what is left of the text + ), + ) + // "word" inside "swordx" is not a whole word, so the shorter "wo" matches instead + assertContentEquals(intArrayOf('s'.code, 6, 'r'.code, 'd'.code, 'x'.code), splitter.encode("swordx")) + assertContentEquals(intArrayOf('a'.code, ' '.code, 5), splitter.encode("a word")) + } + @Test fun `consecutive specials emit no spurious base segments`() { val specials = mapOf("" to 1, "" to 2)