fix(io,cpu): infer BPE for legacy tokenizer.json (#858); gather with N-D indices (#859)

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) {

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<FP32, Float>(
+            Shape(10, 4), FP32::class, FloatArray(40) { it.toFloat() },
+        )
+        val ids = ctx.fromIntArray<Int32, Int>(Shape(2, 3), Int32::class, intArrayOf(0, 1, 2, 7, 8, 9))
+
+        @Suppress("UNCHECKED_CAST")
+        val out = ctx.ops.gather(table, ids as Tensor<sk.ainet.lang.types.DType, *>, 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(),
+        )
+    }
 }

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)

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<UnsupportedTokenizerException> {
+            TokenizerFactory.fromTokenizerJson(json)
+        }
+    }
+
     @Test
     fun `tokenizer_json Unigram dispatches to SentencePiece`() {
         val json = """