From ccbf7e1c6100116b964b162702a536a6c7566d92 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Fri, 24 Jul 2026 13:34:43 +0200 Subject: [PATCH] fix(io,cpu): infer BPE for legacy tokenizer.json; gather with N-D indices MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - TokenizerFactory.fromTokenizerJson (#858): legacy tokenizer.json files that omit model.type (e.g. openai-community/gpt2) are now supported by inferring the type from structure — a merges list is unique to BPE, so such files route to QwenByteLevelBpeTokenizer instead of throwing. - CPU gather (#859): multi-dimensional [N, L] indices threw because a flat data[i] access needs one coordinate per dimension. Now reads the indices in row-major order via the contiguous buffer, falling back to unravel. Adds tests: legacy no-model.type tokenizer.json infers BPE (and still throws without merges); gather with [2,3] indices returns [2,3,4]. Closes #858, #859 --- .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 23 +++++++++++++-- .../exec/tensor/ops/GatherRowDequantTest.kt | 24 +++++++++++++++ .../sk/ainet/io/tokenizer/TokenizerFactory.kt | 25 ++++++++++++++-- .../tokenizer/TokenizerFactoryDispatchTest.kt | 29 +++++++++++++++++++ 4 files changed, 96 insertions(+), 5 deletions(-) diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 5bdd0d04..cd3e9d04 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -2708,9 +2708,26 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory // Input: [vocabSize, embeddingDim], Indices: [L] or [N, L] // Output: [L, embeddingDim] or [N, L, embeddingDim] val numIndices = indices.volume - val indexList = IntArray(numIndices) { i -> - val v = indices.data[i] - (v as Number).toInt() + // Read the indices in row-major order. The vararg element accessor + // requires one coordinate per dimension, so a flat `data[i]` throws for + // rank >= 2 indices — read the contiguous buffer when present, otherwise + // unravel the flat position into a coordinate. + val idxData = indices.data + val indexList = when (idxData) { + is IntArrayTensorData<*> -> IntArray(numIndices) { idxData.buffer[it] } + is FloatArrayTensorData<*> -> IntArray(numIndices) { idxData.buffer[it].toInt() } + else -> { + val dims = indices.shape.dimensions + IntArray(numIndices) { flat -> + val coord = IntArray(dims.size) + var rem = flat + for (d in dims.indices.reversed()) { + coord[d] = rem % dims[d] + rem /= dims[d] + } + (idxData.get(*coord) as Number).toInt() + } + } } if (input.rank == 2) { diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/GatherRowDequantTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/GatherRowDequantTest.kt index 80bebefb..244dc806 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/GatherRowDequantTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/GatherRowDequantTest.kt @@ -42,4 +42,28 @@ class GatherRowDequantTest { out.data.copyToFloatArray(), ) } + + @Test + fun gatherAcceptsMultiDimensionalIndices() { + // Regression for #859: [N, L] indices used to throw because a flat + // `indices.data[i]` needs one coordinate per dimension. + val ctx = DirectCpuExecutionContext.create() + val table = ctx.fromFloatArray( + Shape(10, 4), FP32::class, FloatArray(40) { it.toFloat() }, + ) + val ids = ctx.fromIntArray(Shape(2, 3), Int32::class, intArrayOf(0, 1, 2, 7, 8, 9)) + + @Suppress("UNCHECKED_CAST") + val out = ctx.ops.gather(table, ids as Tensor, dim = 0) + + assertEquals(listOf(2, 3, 4), out.shape.dimensions.toList()) + // Row r of the table is [4r, 4r+1, 4r+2, 4r+3]; indices pick rows 0,1,2 / 7,8,9. + assertContentEquals( + floatArrayOf( + 0f, 1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f, 9f, 10f, 11f, + 28f, 29f, 30f, 31f, 32f, 33f, 34f, 35f, 36f, 37f, 38f, 39f, + ), + out.data.copyToFloatArray(), + ) + } } diff --git a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/TokenizerFactory.kt b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/TokenizerFactory.kt index dddb01af..6dbb8a5f 100644 --- a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/TokenizerFactory.kt +++ b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/TokenizerFactory.kt @@ -66,12 +66,18 @@ public object TokenizerFactory { * to [QwenByteLevelBpeTokenizer]; `"Unigram"` (SentencePiece) gets * wrapped in [SpecialTokenSplitter] when its `added_tokens` registry * is non-empty; `"WordPiece"` currently throws. + * + * Legacy `tokenizer.json` files predating the `model.type` field (e.g. the + * official `openai-community/gpt2` tokenizer) are still supported: the model + * type is inferred from the structure — a `model.merges` list is unique to + * BPE, so such files route to [QwenByteLevelBpeTokenizer]. */ @JvmStatic public fun fromTokenizerJson(json: String): Tokenizer { val root = JSON.parseToJsonElement(json).jsonObject - val modelType = root["model"]?.jsonObject?.get("type")?.jsonPrimitive?.content - ?: throw UnsupportedTokenizerException("tokenizer.json has no model.type") + val model = root["model"]?.jsonObject + ?: throw UnsupportedTokenizerException("tokenizer.json has no 'model'") + val modelType = model["type"]?.jsonPrimitive?.content ?: inferModelType(model) return when (modelType) { "BPE" -> QwenByteLevelBpeTokenizer.fromTokenizerJson(root) "Unigram" -> wrapSentencePieceWithSpecialsFromJson( @@ -87,6 +93,21 @@ public object TokenizerFactory { } } + /** + * Infers the tokenizer's model type for legacy `tokenizer.json` files that + * omit `model.type`. A `merges` list is unique to BPE among the HF model + * types (Unigram and WordPiece have none), so its presence identifies a + * byte-level BPE tokenizer such as GPT-2's. + */ + private fun inferModelType(model: JsonObject): String = + if (model["merges"] != null) { + "BPE" + } else { + throw UnsupportedTokenizerException( + "tokenizer.json has no model.type and it could not be inferred (no 'merges' list)" + ) + } + /** * Apply the [SpecialTokenSplitter] decorator to a SentencePiece base * if the GGUF metadata carries any CONTROL (3) or USER_DEFINED (4) diff --git a/skainet-io/skainet-io-core/src/commonTest/kotlin/sk/ainet/io/tokenizer/TokenizerFactoryDispatchTest.kt b/skainet-io/skainet-io-core/src/commonTest/kotlin/sk/ainet/io/tokenizer/TokenizerFactoryDispatchTest.kt index 1789f4d3..d7881112 100644 --- a/skainet-io/skainet-io-core/src/commonTest/kotlin/sk/ainet/io/tokenizer/TokenizerFactoryDispatchTest.kt +++ b/skainet-io/skainet-io-core/src/commonTest/kotlin/sk/ainet/io/tokenizer/TokenizerFactoryDispatchTest.kt @@ -98,6 +98,35 @@ class TokenizerFactoryDispatchTest { assertEquals(listOf(3, 2), ids.toList()) } + @Test + fun `legacy tokenizer_json without model_type infers BPE from merges`() { + // GPT-2's official tokenizer.json predates the model.type field. + val json = """ + { + "version": "1.0", + "added_tokens": [ + {"id": 2, "content": "<|end|>", "special": true} + ], + "pre_tokenizer": {"type": "ByteLevel"}, + "model": { + "vocab": {"a": 0, "b": 1, "<|end|>": 2, "ab": 3}, + "merges": ["a b"] + } + } + """.trimIndent() + val tok = TokenizerFactory.fromTokenizerJson(json) + assertTrue(tok is QwenByteLevelBpeTokenizer) + assertEquals(listOf(3, 2), tok.encode("ab<|end|>").toList()) + } + + @Test + fun `tokenizer_json without model_type and without merges throws`() { + val json = """{"model":{"vocab":{"a":0}}}""" + assertFailsWith { + TokenizerFactory.fromTokenizerJson(json) + } + } + @Test fun `tokenizer_json Unigram dispatches to SentencePiece`() { val json = """