feat(io-core): SpecialTokenSplitter decorator + SentencePiece HF JSON gaps

skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/SentencePieceTokenizer.kt

@@ -1,6 +1,8 @@
 package sk.ainet.io.tokenizer
 
 import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.boolean
+import kotlinx.serialization.json.booleanOrNull
 import kotlinx.serialization.json.doubleOrNull
 import kotlinx.serialization.json.float
 import kotlinx.serialization.json.int
@@ -319,6 +321,21 @@ public class SentencePieceTokenizer(
          *
          * HF Unigram stores the vocab as a JSON array of `[token, score]`
          * pairs, indexed by id. The unknown token id is at `model.unk_id`.
+         *
+         * `addSpacePrefix` is detected from the `normalizer` block — a
+         * `Prepend` step with `prepend == "▁"` (or a normalizer Sequence
+         * containing one) means the model expects the SP whitespace
+         * dummy prefix. Gemma 4 omits it (HF normalizer is `Replace`-only)
+         * and the model is trained accordingly; without honoring this
+         * flag the encoder produces leading-space tokens (`▁Hi` instead
+         * of `Hi`) that diverge from HuggingFace reference IDs.
+         *
+         * `bosTokenId` / `eosTokenId` are resolved from `added_tokens` by
+         * matching common chat-template names (`<bos>` / `<|begin_of_text|>` /
+         * `<s>` for BOS; `<eos>` / `<|end_of_text|>` / `</s>` for EOS).
+         * HF tokenizer.json files do not carry a dedicated bos/eos field,
+         * so this name-based heuristic mirrors how `transformers` itself
+         * looks them up at load time.
          */
         public fun fromTokenizerJson(root: JsonObject): SentencePieceTokenizer {
             val model = root["model"]?.jsonObject
@@ -336,13 +353,62 @@ public class SentencePieceTokenizer(
             }
 
             val unknownId = model["unk_id"]?.jsonPrimitive?.int
+            val addSpacePrefix = detectAddSpacePrefix(root)
+            val (bosId, eosId) = extractBosEosFromAddedTokens(root)
+
             return SentencePieceTokenizer(
                 tokens = tokens,
                 scores = scores,
                 unknownTokenId = unknownId,
+                bosTokenId = bosId,
+                eosTokenId = eosId,
+                addSpacePrefix = addSpacePrefix,
             )
         }
 
+        /**
+         * Inspect `tokenizer.json#normalizer` for an SP-style `Prepend`
+         * step with `prepend == "▁"`. Returns true when found, false when
+         * a normalizer block exists but no such step is present, true as
+         * a default when there's no normalizer (matches the SPM library
+         * default of `add_dummy_prefix=true`).
+         */
+        private fun detectAddSpacePrefix(root: JsonObject): Boolean {
+            val normalizer = root["normalizer"] as? JsonObject ?: return true
+            val seq = normalizer["normalizers"]?.jsonArray
+            val candidates = seq ?: listOf(normalizer)
+            for (n in candidates) {
+                val obj = n as? JsonObject ?: continue
+                val type = obj["type"]?.jsonPrimitive?.content ?: continue
+                if (type == "Prepend") {
+                    val prepend = obj["prepend"]?.jsonPrimitive?.content
+                    if (prepend == "▁") return true
+                }
+            }
+            return false
+        }
+
+        /**
+         * Resolve BOS/EOS ids from `added_tokens` by matching standard
+         * chat-template content names. First match wins; returns null
+         * for either when no entry matches.
+         */
+        private fun extractBosEosFromAddedTokens(root: JsonObject): Pair<Int?, Int?> {
+            val added = root["added_tokens"]?.jsonArray ?: return null to null
+            var bosId: Int? = null
+            var eosId: Int? = null
+            for (entry in added) {
+                val obj = entry as? JsonObject ?: continue
+                val content = obj["content"]?.jsonPrimitive?.content ?: continue
+                val id = obj["id"]?.jsonPrimitive?.int ?: continue
+                when (content) {
+                    "<bos>", "<|begin_of_text|>", "<s>" -> if (bosId == null) bosId = id
+                    "<eos>", "<|end_of_text|>", "</s>" -> if (eosId == null) eosId = id
+                }
+            }
+            return bosId to eosId
+        }
+
         private const val TOKEN_TYPE_UNKNOWN = 2
     }
 }

skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/SpecialTokenSplitter.kt

@@ -0,0 +1,119 @@
+package sk.ainet.io.tokenizer
+
+/**
+ * Decorator that adds atomic special-token splitting on top of any base
+ * [Tokenizer] that does not already implement it.
+ *
+ * Why this exists
+ * ---------------
+ * Chat-template models embed control markers — `<bos>`, `<|im_start|>`,
+ * `<|turn>`, `<tool_call>`, etc. — that the model is trained to see as a
+ * single id rather than as the byte/SP-encoded fragments their literal
+ * string would otherwise produce. [QwenByteLevelBpeTokenizer] and
+ * [TekkenTokenizer] handle this internally because their input formats
+ * (GGUF, tekken.json) carry an explicit special-token registry.
+ * [SentencePieceTokenizer], by design, does not — it implements pure
+ * llama.cpp-style SPM with byte fallback and nothing else.
+ *
+ * That left a gap for HuggingFace-flavored SentencePiece models like
+ * Gemma 4 whose `tokenizer.json#added_tokens` registry holds the chat-
+ * template specials. This decorator closes the gap: callers (typically
+ * [TokenizerFactory]) construct the bare base tokenizer, extract the
+ * specials map from the source format, and wrap.
+ *
+ * The decorator is intentionally generic — a future refactor can lift
+ * the inline special-token logic out of [QwenByteLevelBpeTokenizer] and
+ * [TekkenTokenizer] and use this decorator there too. For this change
+ * we only apply it to SentencePiece, which is the smallest blast-radius
+ * fix that closes the actual user-visible bug.
+ *
+ * Algorithm
+ * ---------
+ * - **encode(text)**: walk left-to-right; at each position try the
+ *   longest registered special-token string. On a match, emit its id
+ *   and advance past it. Otherwise extend the current non-special
+ *   segment until the next special boundary (or end-of-text), then
+ *   `base.encode(segment)`.
+ * - **decode(ids)**: scan ids; collect contiguous non-special ids into
+ *   a buffer, flushing via `base.decode(buffer)` when we hit a special
+ *   id, then emit the special's string form. The byte-level UTF-8
+ *   spanning that some bases (notably SentencePiece) do inside their
+ *   `decode(IntArray)` is preserved within each non-special run,
+ *   because special-token boundaries always sit on UTF-8 boundaries
+ *   (specials are literal strings).
+ */
+public class SpecialTokenSplitter(
+    private val base: Tokenizer,
+    private val specialTokens: Map<String, Int>,
+    override val bosTokenId: Int? = base.bosTokenId,
+    override val eosTokenId: Int? = base.eosTokenId,
+) : Tokenizer {
+
+    private val specialIdToString: Map<Int, String> =
+        specialTokens.entries.associate { (k, v) -> v to k }
+
+    /** Longest-first ordering so e.g. `<|im_start|>` wins over `<|im`. */
+    private val specialsByLengthDesc: List<String> =
+        specialTokens.keys.sortedByDescending { it.length }
+
+    override val vocabSize: Int = base.vocabSize
+
+    override fun encode(text: String): IntArray {
+        if (specialTokens.isEmpty()) return base.encode(text)
+        val out = ArrayList<Int>(text.length)
+        var i = 0
+        while (i < text.length) {
+            val matched = matchSpecialAt(text, i)
+            if (matched != null) {
+                out.add(specialTokens.getValue(matched))
+                i += matched.length
+                continue
+            }
+            val nextSpecial = nextSpecialStart(text, i)
+            val segment = text.substring(i, nextSpecial)
+            for (id in base.encode(segment)) out.add(id)
+            i = nextSpecial
+        }
+        return IntArray(out.size) { out[it] }
+    }
+
+    override fun decode(ids: IntArray): String {
+        if (ids.isEmpty()) return ""
+        if (specialTokens.isEmpty()) return base.decode(ids)
+        val sb = StringBuilder()
+        val buffer = ArrayList<Int>()
+        for (id in ids) {
+            val special = specialIdToString[id]
+            if (special != null) {
+                if (buffer.isNotEmpty()) {
+                    sb.append(base.decode(buffer.toIntArray()))
+                    buffer.clear()
+                }
+                sb.append(special)
+            } else {
+                buffer.add(id)
+            }
+        }
+        if (buffer.isNotEmpty()) {
+            sb.append(base.decode(buffer.toIntArray()))
+        }
+        return sb.toString()
+    }
+
+    private fun matchSpecialAt(text: String, from: Int): String? {
+        for (tok in specialsByLengthDesc) {
+            if (tok.isNotEmpty() && text.regionMatches(from, tok, 0, tok.length)) return tok
+        }
+        return null
+    }
+
+    private fun nextSpecialStart(text: String, from: Int): Int {
+        var earliest = text.length
+        for (tok in specialTokens.keys) {
+            if (tok.isEmpty()) continue
+            val idx = text.indexOf(tok, startIndex = from + 1)
+            if (idx in 0 until earliest) earliest = idx
+        }
+        return earliest
+    }
+}

skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/TokenizerFactory.kt

@@ -1,6 +1,10 @@
 package sk.ainet.io.tokenizer
 
 import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.boolean
+import kotlinx.serialization.json.int
+import kotlinx.serialization.json.jsonArray
 import kotlinx.serialization.json.jsonObject
 import kotlinx.serialization.json.jsonPrimitive
 import kotlin.jvm.JvmStatic
@@ -42,7 +46,10 @@ public object TokenizerFactory {
             )
         return when (model) {
             "gpt2", "bpe" -> QwenByteLevelBpeTokenizer.fromGgufFields(fields)
-            "llama", "sentencepiece" -> SentencePieceTokenizer.fromGgufFields(fields)
+            "llama", "sentencepiece" -> wrapSentencePieceWithSpecialsFromGguf(
+                base = SentencePieceTokenizer.fromGgufFields(fields),
+                fields = fields,
+            )
             "bert", "wordpiece" -> throw UnsupportedTokenizerException(
                 "WordPiece/BERT tokenizer not yet implemented"
             )
@@ -56,8 +63,9 @@ public object TokenizerFactory {
      * Build a tokenizer from a HuggingFace `tokenizer.json` string.
      *
      * Dispatches on `model.type`: `"BPE"` + byte-level pretokenizer routes
-     * to [QwenByteLevelBpeTokenizer]; `"Unigram"` (SentencePiece) and
-     * `"WordPiece"` currently throw.
+     * to [QwenByteLevelBpeTokenizer]; `"Unigram"` (SentencePiece) gets
+     * wrapped in [SpecialTokenSplitter] when its `added_tokens` registry
+     * is non-empty; `"WordPiece"` currently throws.
      */
     @JvmStatic
     public fun fromTokenizerJson(json: String): Tokenizer {
@@ -66,7 +74,10 @@ public object TokenizerFactory {
             ?: throw UnsupportedTokenizerException("tokenizer.json has no model.type")
         return when (modelType) {
             "BPE" -> QwenByteLevelBpeTokenizer.fromTokenizerJson(root)
-            "Unigram" -> SentencePieceTokenizer.fromTokenizerJson(root)
+            "Unigram" -> wrapSentencePieceWithSpecialsFromJson(
+                base = SentencePieceTokenizer.fromTokenizerJson(root),
+                root = root,
+            )
             "WordPiece" -> throw UnsupportedTokenizerException(
                 "WordPiece tokenizer.json not yet implemented"
             )
@@ -76,6 +87,63 @@ public object TokenizerFactory {
         }
     }
 
+    /**
+     * Apply the [SpecialTokenSplitter] decorator to a SentencePiece base
+     * if the GGUF metadata carries any CONTROL (3) or USER_DEFINED (4)
+     * token-type entries. Both are atomic chat-template markers that the
+     * model expects to see as single ids — `<bos>` is typically CONTROL,
+     * `<|tool_call>` and similar app-specific markers are USER_DEFINED.
+     * The bare base is returned when no specials are present (vanilla
+     * LLaMA-style models) so consumers see no behavior change.
+     */
+    @Suppress("UNCHECKED_CAST")
+    private fun wrapSentencePieceWithSpecialsFromGguf(
+        base: SentencePieceTokenizer,
+        fields: Map<String, Any?>,
+    ): Tokenizer {
+        val tokens = (fields["tokenizer.ggml.tokens"] as? List<*>)
+            ?.filterIsInstance<String>().orEmpty()
+        val tokenTypes = (fields["tokenizer.ggml.token_type"] as? List<*>)
+            ?.mapNotNull { (it as? Number)?.toInt() }.orEmpty()
+        if (tokens.isEmpty() || tokenTypes.isEmpty()) return base
+
+        val specials = HashMap<String, Int>()
+        val limit = minOf(tokens.size, tokenTypes.size)
+        for (i in 0 until limit) {
+            val type = tokenTypes[i]
+            if (type == TOKEN_TYPE_CONTROL || type == TOKEN_TYPE_USER_DEFINED) {
+                val tok = tokens[i]
+                if (tok.isNotEmpty()) specials[tok] = i
+            }
+        }
+        return if (specials.isEmpty()) base else SpecialTokenSplitter(base, specials)
+    }
+
+    /**
+     * Apply the [SpecialTokenSplitter] decorator to a SentencePiece base
+     * built from `tokenizer.json` if its `added_tokens` array carries any
+     * `"special": true` (or unset, defaulting to true) entries. Returns
+     * the bare base when the registry is empty.
+     */
+    private fun wrapSentencePieceWithSpecialsFromJson(
+        base: SentencePieceTokenizer,
+        root: JsonObject,
+    ): Tokenizer {
+        val added = root["added_tokens"]?.jsonArray ?: return base
+        val specials = HashMap<String, Int>(added.size)
+        for (entry in added) {
+            val obj = entry as? JsonObject ?: continue
+            val content = obj["content"]?.jsonPrimitive?.content ?: continue
+            val id = obj["id"]?.jsonPrimitive?.int ?: continue
+            val isSpecial = obj["special"]?.jsonPrimitive?.boolean ?: true
+            if (isSpecial) specials[content] = id
+        }
+        return if (specials.isEmpty()) base else SpecialTokenSplitter(base, specials)
+    }
+
+    private const val TOKEN_TYPE_CONTROL = 3
+    private const val TOKEN_TYPE_USER_DEFINED = 4
+
     internal val JSON: Json = Json { ignoreUnknownKeys = true; isLenient = true }
 }