Embedding coverage Phase 1: BGE (CLS pooling, retrieval prefixes)

CHANGELOG.md

@@ -9,6 +9,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- **BGE embedding models** (`BAAI/bge-small-en-v1.5` and siblings) run on the BERT DSL path:
+  - **CLS pooling.** `BertPooling { MEAN, CLS }` on `BertEncoderRuntime` /
+    `createBertEncoderRuntime`; auto-detected from the sentence-transformers
+    `1_Pooling/config.json` (absent file → `MEAN`, unsupported max/sqrt-len modes rejected
+    loudly). Pooling stays outside the traced graph — OPTIMIZED mode and StableHLO export
+    are unaffected.
+  - **Query/document asymmetry.** `EmbeddingModel` gains `embedQuery` / `embedDocument` /
+    `embedDocuments` (defaults delegate to `embed` — additive, existing consumers untouched).
+    `PrefixedEmbeddingModel` + `EmbeddingModelProfiles` apply retrieval instruction prefixes
+    per repo id (E5 `query: `/`passage: `, BGE query instruction); `fromHuggingFace` wires
+    them automatically, `fromSafeTensors` accepts explicit `prefixes`.
+  - **Integer checkpoint buffers no longer break loads.** BGE-style snapshots persist an
+    I64 `embeddings.position_ids` buffer; the interim `FloatSafeTensorsLoader` skips
+    non-float buffers (the index-free encoder never needs them). Drop when the engine's
+    loader gains a tensor filter ([SKaiNET#822](https://github.com/SKaiNET-developers/SKaiNET/issues/822)).
+  - Design + traceable plan: `docs/specs/embedding-model-coverage.md` (E5 multilingual
+    follows in Phase 2 — Unigram tokenizer).
+
 ## [0.36.0] — 2026-07-12
 
 Ships against **SKaiNET engine 0.36.0**. Headline: **BERT is now completely defined on the DSL

docs/specs/embedding-model-coverage.md

@@ -0,0 +1,115 @@
+# Embedding model coverage: BGE + multilingual E5
+
+**Status:** Phase 1 in progress · **Owner:** @michalharakal · **Created:** 2026-07-13
+
+Extend the BERT embedding stack beyond LEAF to the two most-used compact retrieval
+models, establishing the model matrix for the upcoming embedding **performance
+benchmarks** (`llm-performance` will compare models × execution modes on one harness):
+
+| Target model | Dims | Size | Blocking gaps |
+|---|---|---|---|
+| `BAAI/bge-small-en-v1.5` | 384 | 130 MB | CLS pooling; query instruction prefix |
+| `intfloat/multilingual-e5-small` | 384 | 118 MB | Unigram (XLM-R) tokenizer; `query:`/`passage:` prefixes |
+
+Verified facts (HF configs, 2026-07-13): both are plain 12-layer BERTs (hidden 384,
+heads 12, FFN 1536, absolute positions, `type_vocab_size` 2, **no** `2_Dense` head,
+modules `Transformer → Pooling → Normalize`). `bertNetwork()` runs both **unchanged** —
+every gap is peripheral (tokenizer, pooling, text prefixes).
+
+## Design
+
+### D1. Pooling modes (BGE)
+
+`BertEncoderRuntime` hard-codes masked mean pooling. Add:
+
+```kotlin
+public enum class BertPooling { MEAN, CLS }
+```
+
+- Constructor/`createBertEncoderRuntime` parameter, default `MEAN` (source- and
+  behavior-compatible; BCV dumps regenerate for the new signatures).
+- `encode()`: `CLS` takes hidden-state row 0; `MEAN` keeps the existing masked mean.
+  Pooling stays **outside** the traced encoder graph (unchanged export/OPTIMIZED path).
+- Detection: sentence-transformers `1_Pooling/config.json` (`pooling_mode_cls_token`
+  / `pooling_mode_mean_tokens`) parsed next to `BertConfigParser`; both factories wire
+  it through. Absent file → `MEAN` (status quo).
+
+### D2. Query/document asymmetry (E5 requires, BGE recommends)
+
+Retrieval models embed queries and documents differently (instruction prefixes).
+
+- **SPI (additive):** `EmbeddingModel` gains `embedQuery(text)` and
+  `embedDocument(text)` / `embedDocuments(texts)` with defaults delegating to
+  `embed(...)` — existing implementations and consumers are untouched.
+- **`PrefixedEmbeddingModel`** decorator (llm-providers): prepends per-role prefixes;
+  `embed(...)` uses the document role (indexing is the common bulk path).
+- **`EmbeddingModelProfiles`** registry keyed by repo-id prefix:
+  `intfloat/multilingual-e5-*` → `("query: ", "passage: ")`;
+  `BAAI/bge-*-en-*` → `("Represent this sentence for searching relevant passages: ", "")`.
+  `fromHuggingFace` applies the profile automatically; `fromSafeTensors` accepts
+  explicit prefix parameters (no repo id to key on).
+
+### D3. Unigram tokenizer (E5) — Phase 2
+
+`HuggingFaceTokenizer` is WordPiece-only with hard-coded `[CLS]`/`[SEP]`. E5 needs
+SentencePiece **Unigram** (Viterbi over a 250k scored vocab) plus XLM-R special tokens
+(`<s>`/`</s>`, ids 0/2).
+
+- Source of truth: `tokenizer.json` (embeds the scored vocab + post-processor —
+  no protobuf `.model` parsing).
+- Normalizer: XLM-R uses a `Precompiled` charsmap; we approximate with NFKC and
+  **quantify drift** with a multilingual parity corpus vs HF `tokenizers` (CP-3).
+- Placement decision at CP-3 (see checkpoints): transformers-side
+  (`llm-inference/bert`, no engine release needed) vs engine
+  (`sk.ainet.io.tokenizer`, where SentencePiece-BPE lives, requires engine release).
+  Default: transformers-side first, upstream to the engine later.
+
+### Out of scope
+
+Max/sqrt-len pooling modes, two-segment (cross-encoder) inputs, matryoshka dims,
+serving-side batching.
+
+## Traceable implementation plan
+
+Gitflow: `feature/*` → `develop` → `release/x.y.z` → tag (publishes to Maven Central).
+
+| ID | Work item | Repo / branch | Release | Status |
+|---|---|---|---|---|
+| E5BGE-0 | This design doc | SK-tr `feature/embedding-pooling-profiles` | 0.37.0 | in progress |
+| E5BGE-1 | `BertPooling` (MEAN/CLS) + `1_Pooling` detection + tests | SK-tr, same branch | 0.37.0 | in progress |
+| E5BGE-2 | SPI `embedQuery`/`embedDocument` + `PrefixedEmbeddingModel` + profiles | SK-tr, same branch | 0.37.0 | in progress |
+| E5BGE-3 | BGE end-to-end verification (CLS, 384 dims, prefix) + smoke entry | SK-tr, same branch | 0.37.0 | pending |
+| E5BGE-4 | Unigram tokenizer spike → placement decision | SK-tr (or engine per CP-3) | 0.38.0 | pending |
+| E5BGE-5 | Unigram impl + multilingual parity + special-token generalization | per CP-3 | 0.38.0 | pending |
+| E5BGE-6 | E5 end-to-end verification + docs + benchmark matrix entry | SK-tr | 0.38.0 | pending |
+| E5BGE-7 | leaf-cli: `embedQuery`/`embedDocument` wiring + dim recorded in index | SK-leaf `feature/*` → `develop` | after 0.37.0 | pending |
+
+## Checkpoints (upstream-change gates)
+
+- **CP-1 — Phase 1 code complete** (E5BGE-1..2): full build + `apiCheck` green,
+  pooling/prefix unit tests pass. *Engine changes required: none.*
+- **CP-2 — BGE verified** (E5BGE-3): `fromHuggingFace("BAAI/bge-small-en-v1.5")`
+  produces 384-dim CLS-pooled embeddings; parity vs sentence-transformers reference
+  vectors; PR to `develop`; ship as **0.37.0**. *Gate: confirm no engine release is on
+  the critical path — if anything surfaced, stop and file engine issues first.*
+  **Gate outcome (2026-07-14):** one engine gap surfaced — `SafeTensorsParametersLoader`
+  cannot skip BGE's persisted I64 `embeddings.position_ids` buffer
+  ([SKaiNET#822](https://github.com/SKaiNET-developers/SKaiNET/issues/822)). Bridged by
+  the interim `FloatSafeTensorsLoader` in llm-providers (permute-handler pattern from
+  0.36.0: local, documented, dropped when the engine fix ships). **No engine release on
+  the 0.37.0 critical path.**
+- **CP-3 — Tokenizer placement decision** (E5BGE-4): spike Unigram-from-tokenizer.json;
+  measure normalizer drift on a multilingual corpus. **Decision recorded here**:
+  transformers-side vs engine-side. *This is the explicit "are upstream-engine changes
+  required?" gate for Phase 2.*
+- **CP-4 — E5 verified** (E5BGE-5..6): parity ≤ 1e-5 vs sentence-transformers on the
+  parity corpus (drift from the normalizer approximation documented), smoke green;
+  ship as **0.38.0**.
+
+## Benchmark baseline hook
+
+Once CP-2 lands, the embedding benchmark matrix becomes
+`{mdbr-leaf-ir, bge-small-en-v1.5} × {DIRECT, OPTIMIZED}`, extended by
+`multilingual-e5-small` at CP-4 — measured with the existing `llm-performance`
+harness plus the `PhaseProfile` diagnostic (branch `perf/llama-packed-load-memory`)
+for the non-matmul tail.

llm-api/api/jvm/llm-api.api

@@ -117,6 +117,9 @@ public abstract interface class sk/ainet/llm/api/EmbeddingModel : java/lang/Auto
 	public fun close ()V
 	public fun embed (Ljava/lang/String;)[F
 	public fun embed (Ljava/util/List;)Ljava/util/List;
+	public fun embedDocument (Ljava/lang/String;)[F
+	public fun embedDocuments (Ljava/util/List;)Ljava/util/List;
+	public fun embedQuery (Ljava/lang/String;)[F
 	public abstract fun getDimensions ()I
 }
 

llm-api/src/commonMain/kotlin/sk/ainet/llm/api/EmbeddingModel.kt

@@ -20,6 +20,20 @@ public interface EmbeddingModel : AutoCloseable {
     public fun embed(texts: List<String>): List<FloatArray> =
         call(EmbeddingRequest(texts)).embeddings.sortedBy { it.index }.map { it.vector }
 
+    /**
+     * Embed a *search query*. Retrieval-tuned models embed queries and
+     * documents asymmetrically (instruction prefixes — E5's `query: `, BGE's
+     * `Represent this sentence for searching relevant passages: `); models
+     * without that distinction inherit this default, which is plain [embed].
+     */
+    public fun embedQuery(text: String): FloatArray = embed(text)
+
+    /** Embed a *document / passage* for indexing. See [embedQuery]. */
+    public fun embedDocument(text: String): FloatArray = embed(text)
+
+    /** Batch [embedDocument], vectors in input order. */
+    public fun embedDocuments(texts: List<String>): List<FloatArray> = embed(texts)
+
     /** Output vector dimensionality. */
     public val dimensions: Int
 

llm-inference/bert/api/jvm/bert.api

@@ -6,6 +6,7 @@ public final class sk/ainet/models/bert/BertConfigParser {
 	public static final field INSTANCE Lsk/ainet/models/bert/BertConfigParser;
 	public final fun parse (Ljava/lang/String;Ljava/lang/String;)Lsk/ainet/models/bert/BertModelConfig;
 	public static synthetic fun parse$default (Lsk/ainet/models/bert/BertConfigParser;Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/models/bert/BertModelConfig;
+	public final fun parsePooling (Ljava/lang/String;)Lsk/ainet/models/bert/BertPooling;
 }
 
 public final class sk/ainet/models/bert/BertEmbeddings : sk/ainet/lang/nn/Module, sk/ainet/lang/nn/topology/ModuleParameters {
@@ -17,8 +18,8 @@ public final class sk/ainet/models/bert/BertEmbeddings : sk/ainet/lang/nn/Module
 }
 
 public final class sk/ainet/models/bert/BertEncoderRuntime {
-	public fun <init> (Lsk/ainet/lang/nn/Module;Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/models/bert/BertExecutionMode;)V
-	public synthetic fun <init> (Lsk/ainet/lang/nn/Module;Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/models/bert/BertExecutionMode;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (Lsk/ainet/lang/nn/Module;Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/models/bert/BertExecutionMode;Lsk/ainet/models/bert/BertPooling;)V
+	public synthetic fun <init> (Lsk/ainet/lang/nn/Module;Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/models/bert/BertExecutionMode;Lsk/ainet/models/bert/BertPooling;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun compile (I)Ljava/util/List;
 	public static synthetic fun compile$default (Lsk/ainet/models/bert/BertEncoderRuntime;IILjava/lang/Object;)Ljava/util/List;
 	public final fun encode ([I[I)[F
@@ -73,6 +74,14 @@ public final class sk/ainet/models/bert/BertNetworkLoader {
 	public final fun normalizeTensorNames (Ljava/util/List;)Ljava/util/List;
 }
 
+public final class sk/ainet/models/bert/BertPooling : java/lang/Enum {
+	public static final field CLS Lsk/ainet/models/bert/BertPooling;
+	public static final field MEAN Lsk/ainet/models/bert/BertPooling;
+	public static fun getEntries ()Lkotlin/enums/EnumEntries;
+	public static fun valueOf (Ljava/lang/String;)Lsk/ainet/models/bert/BertPooling;
+	public static fun values ()[Lsk/ainet/models/bert/BertPooling;
+}
+
 public final class sk/ainet/models/bert/HuggingFaceTokenizer : sk/ainet/apps/llm/Tokenizer {
 	public static final field Companion Lsk/ainet/models/bert/HuggingFaceTokenizer$Companion;
 	public fun decode (I)Ljava/lang/String;

llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertConfig.kt

@@ -67,4 +67,23 @@ public object BertConfigParser {
             projectionDim = projectionDim,
         )
     }
+
+    /**
+     * Parses a sentence-transformers `1_Pooling/config.json` into a
+     * [BertPooling] mode. `pooling_mode_cls_token: true` selects
+     * [BertPooling.CLS]; everything else (including `null` for snapshots
+     * without the file) falls back to [BertPooling.MEAN] — the pre-existing
+     * behavior. Max and sqrt-len pooling are rejected explicitly rather than
+     * silently mis-pooled.
+     */
+    public fun parsePooling(poolingConfigJson: String?): BertPooling {
+        if (poolingConfigJson == null) return BertPooling.MEAN
+        fun flag(key: String) =
+            Regex("\"$key\"\\s*:\\s*(true|false)").find(poolingConfigJson)?.groupValues?.get(1) == "true"
+
+        require(!flag("pooling_mode_max_tokens") && !flag("pooling_mode_mean_sqrt_len_tokens")) {
+            "Unsupported pooling mode in 1_Pooling/config.json (max / mean_sqrt_len)"
+        }
+        return if (flag("pooling_mode_cls_token")) BertPooling.CLS else BertPooling.MEAN
+    }
 }

llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertEncoderRuntime.kt

@@ -19,6 +19,7 @@ import sk.ainet.lang.tensor.div
 import sk.ainet.lang.tensor.matmul
 import sk.ainet.lang.tensor.mean
 import sk.ainet.lang.tensor.plus
+import sk.ainet.lang.tensor.reshape
 import sk.ainet.lang.tensor.sqrt
 import sk.ainet.lang.tensor.sum
 import sk.ainet.lang.tensor.t
@@ -36,6 +37,20 @@ public enum class BertExecutionMode {
     OPTIMIZED,
 }
 
+/**
+ * Sentence-pooling strategy applied by [BertEncoderRuntime.encode] on top of
+ * the encoder's hidden states — the sentence-transformers `1_Pooling` module.
+ * Pooling stays outside the traced graph, so the choice has no effect on the
+ * OPTIMIZED path or the StableHLO export.
+ */
+public enum class BertPooling {
+    /** Mask-weighted mean over token positions (LEAF, E5, MiniLM). */
+    MEAN,
+
+    /** Hidden state of the first token — `[CLS]` (BGE, GTE). */
+    CLS,
+}
+
 /**
  * Encoder runtime for BERT sentence embeddings on the DSL path.
  *
@@ -61,6 +76,7 @@ public class BertEncoderRuntime<T : DType>(
     private val projectionWeight: Tensor<T, Float>? = null,
     private val projectionBias: Tensor<T, Float>? = null,
     private val mode: BertExecutionMode = BertExecutionMode.DIRECT,
+    private val pooling: BertPooling = BertPooling.MEAN,
 ) {
 
     /** Output dimensionality of [encode]: projection out-features when present, else hidden size. */
@@ -79,32 +95,37 @@ public class BertEncoderRuntime<T : DType>(
     }
 
     /**
-     * Encode tokens into a single embedding vector: forward → mean pooling
-     * (mask-weighted when [attentionMask] is given) → optional dense
-     * projection → L2 normalization.
+     * Encode tokens into a single embedding vector: forward → pooling
+     * ([BertPooling.MEAN], mask-weighted when [attentionMask] is given, or
+     * [BertPooling.CLS]) → optional dense projection → L2 normalization.
      *
      * @param tokenIds token IDs including `[CLS]` and `[SEP]`
-     * @param attentionMask 1 for real tokens, 0 for padding; affects pooling only
+     * @param attentionMask 1 for real tokens, 0 for padding; affects MEAN pooling only
      * @return normalized vector of size [dimensions]
      */
     public fun encode(tokenIds: IntArray, attentionMask: IntArray? = null): FloatArray {
         val hiddenStates = forward(tokenIds)
         val seqLen = tokenIds.size
 
-        var pooled = if (attentionMask != null) {
-            require(attentionMask.size == seqLen) {
-                "BertEncoderRuntime: attentionMask size ${attentionMask.size} != tokenIds size $seqLen"
+        var pooled = when (pooling) {
+            BertPooling.CLS ->
+                // First row of [L, hidden] — the [CLS] position — as a [hidden] vector.
+                ctx.ops.narrow(hiddenStates, 0, 0, 1).reshape(Shape(config.hiddenSize))
+            BertPooling.MEAN -> if (attentionMask != null) {
+                require(attentionMask.size == seqLen) {
+                    "BertEncoderRuntime: attentionMask size ${attentionMask.size} != tokenIds size $seqLen"
+                }
+                val maskTensor = ctx.fromFloatArray<T, Float>(
+                    Shape(seqLen, 1), dtype,
+                    FloatArray(seqLen) { attentionMask[it].toFloat() }
+                )
+                val masked = hiddenStates * maskTensor
+                val summed = masked.sum(dim = 0)
+                val count = attentionMask.sumOf { it }.toFloat().coerceAtLeast(1f)
+                summed / count
+            } else {
+                hiddenStates.mean(dim = 0)
             }
-            val maskTensor = ctx.fromFloatArray<T, Float>(
-                Shape(seqLen, 1), dtype,
-                FloatArray(seqLen) { attentionMask[it].toFloat() }
-            )
-            val masked = hiddenStates * maskTensor
-            val summed = masked.sum(dim = 0)
-            val count = attentionMask.sumOf { it }.toFloat().coerceAtLeast(1f)
-            summed / count
-        } else {
-            hiddenStates.mean(dim = 0)
         }
 
         if (projectionWeight != null) {
@@ -295,6 +316,7 @@ public inline fun <reified T : DType> createBertEncoderRuntime(
     resolver: WeightNameResolver = BertSafeTensorsNameResolver(),
     mode: BertExecutionMode = BertExecutionMode.DIRECT,
     debug: Boolean = false,
+    pooling: BertPooling = BertPooling.MEAN,
 ): BertEncoderRuntime<T> {
     val model = bertNetwork<T, Float>(config)
 
@@ -337,5 +359,6 @@ public inline fun <reified T : DType> createBertEncoderRuntime(
         projectionWeight = projectionWeight,
         projectionBias = projectionBias,
         mode = mode,
+        pooling = pooling,
     )
 }

llm-inference/bert/src/commonTest/kotlin/sk/ainet/models/bert/BertPoolingConfigTest.kt

@@ -0,0 +1,52 @@
+package sk.ainet.models.bert
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+
+/** [BertConfigParser.parsePooling] against real sentence-transformers `1_Pooling/config.json` shapes. */
+class BertPoolingConfigTest {
+
+    @Test
+    fun clsPooling_bgeShape() {
+        val json = """
+            {
+              "word_embedding_dimension": 384,
+              "pooling_mode_cls_token": true,
+              "pooling_mode_mean_tokens": false,
+              "pooling_mode_max_tokens": false,
+              "pooling_mode_mean_sqrt_len_tokens": false
+            }
+        """.trimIndent()
+        assertEquals(BertPooling.CLS, BertConfigParser.parsePooling(json))
+    }
+
+    @Test
+    fun meanPooling_e5Shape() {
+        val json = """
+            {
+              "word_embedding_dimension": 384,
+              "pooling_mode_cls_token": false,
+              "pooling_mode_mean_tokens": true,
+              "pooling_mode_max_tokens": false,
+              "pooling_mode_mean_sqrt_len_tokens": false
+            }
+        """.trimIndent()
+        assertEquals(BertPooling.MEAN, BertConfigParser.parsePooling(json))
+    }
+
+    @Test
+    fun missingFile_defaultsToMean() {
+        assertEquals(BertPooling.MEAN, BertConfigParser.parsePooling(null))
+    }
+
+    @Test
+    fun unsupportedModes_areRejectedNotSilentlyMispooled() {
+        assertFailsWith<IllegalArgumentException> {
+            BertConfigParser.parsePooling("""{"pooling_mode_max_tokens": true}""")
+        }
+        assertFailsWith<IllegalArgumentException> {
+            BertConfigParser.parsePooling("""{"pooling_mode_mean_sqrt_len_tokens": true}""")
+        }
+    }
+}