diff --git a/CHANGELOG.md b/CHANGELOG.md index a4dd6a35..7ac2e76d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### 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. - **`CoroutineSchedule.forRange` could call its body with an empty, inverted range** (`skainet-backend-cpu`): the chunk length is rounded up, so for some `n` and task counts the last chunk started past `n` (63 elements on 12 tasks: chunks of 6, the twelfth at 66..63). The engine's own bodies loop over the range and did nothing for it; a body that passes the range length to a kernel saw a negative size. The task count now follows the chunk length, and every range is non-empty and covers `[0, n)` as the `Schedule` contract states. ### Changed 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)