feat(io): NameMap — GGUF/HF tensor names ↔ TensorId for Llama, Qwen2, Gemma-3; unmapped reported (SKEEP-003 P1, S0.6)

skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/weights/NameMap.kt

@@ -0,0 +1,109 @@
+package sk.ainet.io.weights
+
+import sk.ainet.lang.tensor.TensorId
+
+/**
+ * Bidirectional mapping between a checkpoint's tensor names (GGUF `blk.3.attn_q.weight`, HF
+ * `model.layers.3.self_attn.q_proj.weight`) and SKaiNET [TensorId]s (SKEEP-003 §4.7 rule 5).
+ * One per checkpoint format and model family; nothing is dropped silently — [unmapped] lists what
+ * a map does not know.
+ *
+ * The canonical SKaiNET ids for transformer checkpoints are family-neutral and GGUF-role based:
+ * `model.embed_tokens.weight`, `model.norm.weight`, `model.lm_head.weight`, `model.rope_freqs.weight`,
+ * and per layer `model.layers[N].attn.{q,k,v,o}_proj.{weight,bias}`, `model.layers[N].attn.{q,k}_norm.weight`,
+ * `model.layers[N].mlp.{gate,up,down}_proj.weight`, `model.layers[N].{attn_norm,ffn_norm,post_attention_norm,post_ffw_norm}.weight`.
+ */
+public interface NameMap {
+    /** Model family this map is for (`llama`, `qwen2`, `gemma3`). */
+    public val family: String
+
+    /** Checkpoint format this map reads (`gguf`, `safetensors`). */
+    public val format: String
+
+    /** The [TensorId] for a checkpoint tensor name, or `null` if this map does not know the name. */
+    public fun toTensorId(checkpointName: String): TensorId?
+
+    /** The checkpoint tensor name for [id], or `null` if the id has no tensor in this format/family. */
+    public fun toCheckpointName(id: TensorId): String?
+
+    /** Names of [checkpointNames] this map cannot translate — never dropped, always reported. */
+    public fun unmapped(checkpointNames: Iterable<String>): List<String> = checkpointNames.filter { toTensorId(it) == null }
+
+    /** Translate every name; unknown names map to `null` so callers can decide. */
+    public fun toTensorIds(checkpointNames: Iterable<String>): Map<String, TensorId?> = checkpointNames.associateWith { toTensorId(it) }
+
+    /**
+     * This map as a legacy [WeightNameResolver]: `(modulePath, paramName)` in the slash form
+     * (`model/layers[3]/attn`, `q_proj.weight`) → checkpoint name.
+     */
+    public fun asWeightNameResolver(): WeightNameResolver = WeightNameResolver { modulePath, paramName ->
+        val segments = (modulePath.takeIf { it.isNotEmpty() }?.split("/") ?: emptyList())
+        val parts = paramName.split(".")
+        // a dotted paramName ("q_proj.weight") contributes its leading parts to the module path
+        toCheckpointName(TensorId(segments + parts.dropLast(1), parts.last()))
+    }
+}
+
+/** SAM-style constructor for [WeightNameResolver]. */
+public fun WeightNameResolver(block: (modulePath: String, paramName: String) -> String?): WeightNameResolver =
+    object : WeightNameResolver {
+        override fun resolve(modulePath: String, paramName: String): String? = block(modulePath, paramName)
+    }
+
+/**
+ * A [NameMap] built from role tables. Checkpoint names are either top-level (`token_embd.weight`)
+ * or per-layer (`blk.3.attn_q.weight` / `model.layers.3.self_attn.q_proj.weight`); the tables map
+ * the format's role spelling to the canonical id suffix and back.
+ *
+ * @property layerPrefix how a layer index is spelled in the checkpoint, with `{N}` as the placeholder
+ *   (`blk.{N}.` for GGUF, `model.layers.{N}.` for HF)
+ * @property topLevel checkpoint top-level tensor stem → canonical id stem (`token_embd` → `model.embed_tokens`)
+ * @property layerRoles checkpoint per-layer role stem → canonical id suffix (`attn_q` → `attn.q_proj`)
+ * @property suffixes accepted trailing parameter names (`weight`, `bias`)
+ */
+public class RoleTableNameMap(
+    override val family: String,
+    override val format: String,
+    private val layerPrefix: String,
+    private val topLevel: Map<String, String>,
+    private val layerRoles: Map<String, String>,
+    private val suffixes: Set<String> = setOf("weight", "bias"),
+) : NameMap {
+    private val layerRegex: Regex
+    private val topLevelInverse: Map<String, String> = topLevel.entries.associate { (k, v) -> v to k }
+    private val layerRolesInverse: Map<String, String> = layerRoles.entries.associate { (k, v) -> v to k }
+
+    init {
+        require("{N}" in layerPrefix) { "layerPrefix must contain {N}: '$layerPrefix'" }
+        val (before, after) = layerPrefix.split("{N}", limit = 2)
+        layerRegex = Regex("^" + Regex.escape(before) + "(\\d+)" + Regex.escape(after) + "(.+)\\.(" + suffixes.joinToString("|") { Regex.escape(it) } + ")$")
+    }
+
+    override fun toTensorId(checkpointName: String): TensorId? {
+        layerRegex.matchEntire(checkpointName)?.let { m ->
+            val layer = m.groupValues[1].toInt()
+            val role = layerRoles[m.groupValues[2]] ?: return null
+            val suffix = m.groupValues[3]
+            return TensorId(listOf("model", "layers[$layer]") + role.split("."), suffix)
+        }
+        val dot = checkpointName.lastIndexOf('.')
+        if (dot <= 0) return null
+        val stem = checkpointName.substring(0, dot); val suffix = checkpointName.substring(dot + 1)
+        if (suffix !in suffixes) return null
+        val canonicalStem = topLevel[stem] ?: return null
+        val parts = canonicalStem.split(".")
+        return TensorId(parts, suffix)
+    }
+
+    override fun toCheckpointName(id: TensorId): String? {
+        if (id.discriminator != null || id.parameter !in suffixes) return null
+        val path = id.modulePath
+        if (path.size >= 3 && path[0] == "model" && path[1].startsWith("layers[") && path[1].endsWith("]")) {
+            val layer = path[1].removePrefix("layers[").removeSuffix("]").toIntOrNull() ?: return null
+            val roleStem = layerRolesInverse[path.drop(2).joinToString(".")] ?: return null
+            return layerPrefix.replace("{N}", layer.toString()) + roleStem + "." + id.parameter
+        }
+        val stem = topLevelInverse[path.joinToString(".")] ?: return null
+        return "$stem.${id.parameter}"
+    }
+}

skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/weights/TransformerNameMaps.kt

@@ -0,0 +1,75 @@
+package sk.ainet.io.weights
+
+/**
+ * Name maps for the transformer families SKaiNET loads (SKEEP-003 M0-A2 reference models:
+ * Llama-3.2-1B, Qwen2.5-0.5B, Gemma-3-1B). GGUF maps in [Gguf], Hugging Face / SafeTensors maps in
+ * [Hf]. All share the canonical ids documented on [NameMap].
+ */
+public object TransformerNameMaps {
+
+    // Canonical id suffixes shared by every family.
+    private const val ATTN_Q = "attn.q_proj"; private const val ATTN_K = "attn.k_proj"; private const val ATTN_V = "attn.v_proj"; private const val ATTN_O = "attn.o_proj"
+    private const val ATTN_Q_NORM = "attn.q_norm"; private const val ATTN_K_NORM = "attn.k_norm"
+    private const val FFN_GATE = "mlp.gate_proj"; private const val FFN_UP = "mlp.up_proj"; private const val FFN_DOWN = "mlp.down_proj"
+    private const val ATTN_NORM = "attn_norm"; private const val FFN_NORM = "ffn_norm"
+    private const val POST_ATTN_NORM = "post_attention_norm"; private const val POST_FFW_NORM = "post_ffw_norm"
+
+    private val topLevelGguf = mapOf(
+        "token_embd" to "model.embed_tokens",
+        "output_norm" to "model.norm",
+        "output" to "model.lm_head",
+        "rope_freqs" to "model.rope_freqs",
+    )
+    private val topLevelHf = mapOf(
+        "model.embed_tokens" to "model.embed_tokens",
+        "model.norm" to "model.norm",
+        "lm_head" to "model.lm_head",
+    )
+
+    private val ggufLayerRoles = mapOf(
+        "attn_q" to ATTN_Q, "attn_k" to ATTN_K, "attn_v" to ATTN_V, "attn_output" to ATTN_O,
+        "attn_q_norm" to ATTN_Q_NORM, "attn_k_norm" to ATTN_K_NORM,
+        "attn_norm" to ATTN_NORM, "ffn_norm" to FFN_NORM,
+        "post_attention_norm" to POST_ATTN_NORM, "post_ffw_norm" to POST_FFW_NORM,
+        "ffn_gate" to FFN_GATE, "ffn_up" to FFN_UP, "ffn_down" to FFN_DOWN,
+    )
+
+    /** HF Llama / Qwen2: `input_layernorm` is the attention norm, `post_attention_layernorm` the pre-FFN (`ffn_norm`). */
+    private val hfLlamaLayerRoles = mapOf(
+        "self_attn.q_proj" to ATTN_Q, "self_attn.k_proj" to ATTN_K, "self_attn.v_proj" to ATTN_V, "self_attn.o_proj" to ATTN_O,
+        "self_attn.q_norm" to ATTN_Q_NORM, "self_attn.k_norm" to ATTN_K_NORM,
+        "input_layernorm" to ATTN_NORM, "post_attention_layernorm" to FFN_NORM,
+        "mlp.gate_proj" to FFN_GATE, "mlp.up_proj" to FFN_UP, "mlp.down_proj" to FFN_DOWN,
+    )
+
+    /** HF Gemma-3 has four norms per layer; `post_attention_layernorm` is a true post-attention norm here. */
+    private val hfGemma3LayerRoles = mapOf(
+        "self_attn.q_proj" to ATTN_Q, "self_attn.k_proj" to ATTN_K, "self_attn.v_proj" to ATTN_V, "self_attn.o_proj" to ATTN_O,
+        "self_attn.q_norm" to ATTN_Q_NORM, "self_attn.k_norm" to ATTN_K_NORM,
+        "input_layernorm" to ATTN_NORM, "post_attention_layernorm" to POST_ATTN_NORM,
+        "pre_feedforward_layernorm" to FFN_NORM, "post_feedforward_layernorm" to POST_FFW_NORM,
+        "mlp.gate_proj" to FFN_GATE, "mlp.up_proj" to FFN_UP, "mlp.down_proj" to FFN_DOWN,
+    )
+
+    /** GGUF maps (`blk.N.<role>.<weight|bias>`). The role set is the union llama.cpp emits for these families. */
+    public object Gguf {
+        public val llama: NameMap = RoleTableNameMap("llama", "gguf", "blk.{N}.", topLevelGguf, ggufLayerRoles)
+        public val qwen2: NameMap = RoleTableNameMap("qwen2", "gguf", "blk.{N}.", topLevelGguf, ggufLayerRoles)
+        public val gemma3: NameMap = RoleTableNameMap("gemma3", "gguf", "blk.{N}.", topLevelGguf, ggufLayerRoles)
+
+        /** Map for a GGUF `general.architecture` value, or `null` if the family is unknown. */
+        public fun forArchitecture(architecture: String): NameMap? = when (architecture.lowercase()) {
+            "llama", "llama4", "mistral", "smollm", "apertus" -> llama
+            "qwen2", "qwen2.5", "qwen3" -> qwen2
+            "gemma3", "gemma2", "gemma" -> gemma3
+            else -> null
+        }
+    }
+
+    /** Hugging Face / SafeTensors maps (`model.layers.N.<role>.<weight|bias>`). */
+    public object Hf {
+        public val llama: NameMap = RoleTableNameMap("llama", "safetensors", "model.layers.{N}.", topLevelHf, hfLlamaLayerRoles)
+        public val qwen2: NameMap = RoleTableNameMap("qwen2", "safetensors", "model.layers.{N}.", topLevelHf, hfLlamaLayerRoles)
+        public val gemma3: NameMap = RoleTableNameMap("gemma3", "safetensors", "model.layers.{N}.", topLevelHf, hfGemma3LayerRoles)
+    }
+}

skainet-io/skainet-io-core/src/commonTest/kotlin/sk/ainet/io/weights/NameMapTest.kt

@@ -0,0 +1,104 @@
+package sk.ainet.io.weights
+
+import sk.ainet.lang.tensor.TensorId
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNull
+import kotlin.test.assertTrue
+
+/**
+ * SKEEP-003 M0-F4 / M0-A2: every tensor of the three reference GGUFs (Llama-3.2-1B, Qwen2.5-0.5B,
+ * Gemma-3-1B) maps to a TensorId with zero unmapped names; round trips hold; unknown names are
+ * reported, never dropped. The tensor lists below are the names llama.cpp's converter emits for
+ * these models (the real files are checked opportunistically by the fixture-gated io-gguf test).
+ */
+class NameMapTest {
+
+    /** Llama-3.2-1B-Instruct GGUF: 16 layers, tied embeddings (no output.weight), rope_freqs. */
+    private fun llama32_1b(): List<String> = buildList {
+        add("token_embd.weight"); add("output_norm.weight"); add("rope_freqs.weight")
+        for (n in 0 until 16) for (r in listOf("attn_norm", "attn_q", "attn_k", "attn_v", "attn_output", "ffn_norm", "ffn_gate", "ffn_up", "ffn_down")) add("blk.$n.$r.weight")
+    }
+
+    /** Qwen2.5-0.5B-Instruct GGUF: 24 layers, q/k/v biases, tied embeddings. */
+    private fun qwen25_05b(): List<String> = buildList {
+        add("token_embd.weight"); add("output_norm.weight")
+        for (n in 0 until 24) {
+            for (r in listOf("attn_norm", "attn_q", "attn_k", "attn_v", "attn_output", "ffn_norm", "ffn_gate", "ffn_up", "ffn_down")) add("blk.$n.$r.weight")
+            for (r in listOf("attn_q", "attn_k", "attn_v")) add("blk.$n.$r.bias")
+        }
+    }
+
+    /** Gemma-3-1B-it GGUF: 26 layers, q/k norms, four norms per layer, tied embeddings. */
+    private fun gemma3_1b(): List<String> = buildList {
+        add("token_embd.weight"); add("output_norm.weight")
+        for (n in 0 until 26) for (r in listOf(
+            "attn_norm", "attn_q", "attn_k", "attn_v", "attn_output", "attn_q_norm", "attn_k_norm",
+            "post_attention_norm", "ffn_norm", "ffn_gate", "ffn_up", "ffn_down", "post_ffw_norm",
+        )) add("blk.$n.$r.weight")
+    }
+
+    private fun checkFamily(map: NameMap, names: List<String>) {
+        assertEquals(emptyList(), map.unmapped(names), "${map.family}: unmapped")
+        val ids = map.toTensorIds(names)
+        assertEquals(names.size, ids.values.filterNotNull().toSet().size, "${map.family}: ids must be distinct")
+        for (n in names) assertEquals(n, map.toCheckpointName(ids[n]!!), "${map.family}: round trip of $n")
+    }
+
+    @Test fun llamaGgufMapsEveryTensor() = checkFamily(TransformerNameMaps.Gguf.llama, llama32_1b())
+    @Test fun qwen2GgufMapsEveryTensor() = checkFamily(TransformerNameMaps.Gguf.qwen2, qwen25_05b())
+    @Test fun gemma3GgufMapsEveryTensor() = checkFamily(TransformerNameMaps.Gguf.gemma3, gemma3_1b())
+
+    @Test
+    fun canonicalIdsAreFamilyNeutral() {
+        val m = TransformerNameMaps.Gguf.llama
+        assertEquals("model.layers[3].attn.q_proj.weight", m.toTensorId("blk.3.attn_q.weight")!!.canonical)
+        assertEquals("model.layers[3].attn.q_proj.bias", TransformerNameMaps.Gguf.qwen2.toTensorId("blk.3.attn_q.bias")!!.canonical)
+        assertEquals("model.layers[0].post_ffw_norm.weight", TransformerNameMaps.Gguf.gemma3.toTensorId("blk.0.post_ffw_norm.weight")!!.canonical)
+        assertEquals("model.embed_tokens.weight", m.toTensorId("token_embd.weight")!!.canonical)
+        assertEquals("model.norm.weight", m.toTensorId("output_norm.weight")!!.canonical)
+        assertEquals("model.lm_head.weight", m.toTensorId("output.weight")!!.canonical)
+        assertEquals("model.rope_freqs.weight", m.toTensorId("rope_freqs.weight")!!.canonical)
+        // the same id resolves to the HF spelling through the HF map
+        val id = m.toTensorId("blk.3.attn_q.weight")!!
+        assertEquals("model.layers.3.self_attn.q_proj.weight", TransformerNameMaps.Hf.llama.toCheckpointName(id))
+        assertEquals("lm_head.weight", TransformerNameMaps.Hf.llama.toCheckpointName(TensorId.parse("model.lm_head.weight")))
+    }
+
+    @Test
+    fun hfNormsDifferPerFamily() {
+        // Llama/Qwen2: post_attention_layernorm is the pre-FFN norm; Gemma-3 has a true post-attention norm
+        assertEquals("model.layers[2].ffn_norm.weight", TransformerNameMaps.Hf.llama.toTensorId("model.layers.2.post_attention_layernorm.weight")!!.canonical)
+        assertEquals("model.layers[2].post_attention_norm.weight", TransformerNameMaps.Hf.gemma3.toTensorId("model.layers.2.post_attention_layernorm.weight")!!.canonical)
+        assertEquals("model.layers[2].ffn_norm.weight", TransformerNameMaps.Hf.gemma3.toTensorId("model.layers.2.pre_feedforward_layernorm.weight")!!.canonical)
+        assertEquals("model.layers[2].post_ffw_norm.weight", TransformerNameMaps.Hf.gemma3.toTensorId("model.layers.2.post_feedforward_layernorm.weight")!!.canonical)
+        // Gemma-3's GGUF post_attention_norm ↔ HF post_attention_layernorm
+        val g = TransformerNameMaps.Gguf.gemma3.toTensorId("blk.2.post_attention_norm.weight")!!
+        assertEquals("model.layers.2.post_attention_layernorm.weight", TransformerNameMaps.Hf.gemma3.toCheckpointName(g))
+    }
+
+    @Test
+    fun unknownNamesAreReportedNotDropped() {
+        val m = TransformerNameMaps.Gguf.llama
+        val names = listOf("token_embd.weight", "blk.0.attn_q.weight", "blk.0.ffn_gate_exps.weight", "some.unknown", "blk.x.attn_q.weight", "blk.0.attn_q.scale")
+        assertEquals(listOf("blk.0.ffn_gate_exps.weight", "some.unknown", "blk.x.attn_q.weight", "blk.0.attn_q.scale"), m.unmapped(names))
+        assertEquals(6, m.toTensorIds(names).size)
+        assertNull(m.toTensorIds(names)["some.unknown"])
+        // ids this family has no tensor for
+        assertNull(m.toCheckpointName(TensorId.parse("model.layers[0].attn.q_proj.weight#step=1")))
+        assertNull(m.toCheckpointName(TensorId.parse("model.layers[0].mystery.weight")))
+        assertNull(m.toCheckpointName(TensorId.parse("model.layers[a].attn.q_proj.weight")))
+    }
+
+    @Test
+    fun legacyResolverAdapter() {
+        val r = TransformerNameMaps.Gguf.llama.asWeightNameResolver()
+        assertEquals("blk.3.attn_q.weight", r.resolve("model/layers[3]/attn", "q_proj.weight"))
+        assertEquals("blk.3.ffn_norm.weight", r.resolve("model/layers[3]", "ffn_norm.weight"))
+        assertEquals("token_embd.weight", r.resolve("model", "embed_tokens.weight"))
+        assertNull(r.resolve("model/layers[3]/attn", "nope.weight"))
+        assertTrue(TransformerNameMaps.Gguf.forArchitecture("llama") === TransformerNameMaps.Gguf.llama)
+        assertTrue(TransformerNameMaps.Gguf.forArchitecture("Qwen2") === TransformerNameMaps.Gguf.qwen2)
+        assertNull(TransformerNameMaps.Gguf.forArchitecture("rwkv"))
+    }
+}

skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufNameMaps.kt

@@ -0,0 +1,40 @@
+package sk.ainet.io.gguf
+
+import sk.ainet.io.weights.NameMap
+import sk.ainet.io.weights.TransformerNameMaps
+import sk.ainet.lang.tensor.TensorId
+
+/** The [NameMap] for this GGUF's `general.architecture`, or `null` if the family is not known. */
+public fun StreamingGGUFReader.nameMap(): NameMap? =
+    (fields["general.architecture"] as? String)?.let { TransformerNameMaps.Gguf.forArchitecture(it) }
+
+/**
+ * Translate every tensor of this GGUF to a [TensorId] through [map] (default: the map for the
+ * file's architecture). Unknown names are kept with a `null` id — see [NameMap.unmapped].
+ */
+public fun StreamingGGUFReader.tensorIds(map: NameMap? = nameMap()): Map<String, TensorId?> =
+    tensors.associate { it.name to map?.toTensorId(it.name) }
+
+/**
+ * This map as the legacy role-based [TensorNameMapper] (GGUF names by role and layer). Roles the
+ * map cannot produce throw [IllegalStateException] — a `TensorNameMapper` has no "absent" answer.
+ */
+public fun NameMap.asTensorNameMapper(): TensorNameMapper = object : TensorNameMapper {
+    private fun name(id: TensorId): String =
+        toCheckpointName(id) ?: throw IllegalStateException("$family/$format name map has no tensor for $id")
+    private fun layer(n: Int, vararg rest: String, parameter: String = "weight") =
+        name(TensorId(listOf("model", "layers[$n]") + rest.toList(), parameter))
+
+    override fun tokenEmbedding(): String = name(TensorId(listOf("model", "embed_tokens"), "weight"))
+    override fun outputNorm(): String = name(TensorId(listOf("model", "norm"), "weight"))
+    override fun outputWeight(): String = name(TensorId(listOf("model", "lm_head"), "weight"))
+    override fun layerAttnNorm(layer: Int): String = layer(layer, "attn_norm")
+    override fun layerAttnQ(layer: Int): String = layer(layer, "attn", "q_proj")
+    override fun layerAttnK(layer: Int): String = layer(layer, "attn", "k_proj")
+    override fun layerAttnV(layer: Int): String = layer(layer, "attn", "v_proj")
+    override fun layerAttnO(layer: Int): String = layer(layer, "attn", "o_proj")
+    override fun layerFfnNorm(layer: Int): String = layer(layer, "ffn_norm")
+    override fun layerFfnGate(layer: Int): String = layer(layer, "mlp", "gate_proj")
+    override fun layerFfnUp(layer: Int): String = layer(layer, "mlp", "up_proj")
+    override fun layerFfnDown(layer: Int): String = layer(layer, "mlp", "down_proj")
+}