diff --git a/CHANGELOG.md b/CHANGELOG.md index dd025d3f..db0aa05e 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/specs/embedding-model-coverage.md b/docs/specs/embedding-model-coverage.md new file mode 100644 index 00000000..f25bbe04 --- /dev/null +++ b/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 +(``/``, 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. diff --git a/llm-api/api/jvm/llm-api.api b/llm-api/api/jvm/llm-api.api index 3207bc29..271be6c8 100644 --- a/llm-api/api/jvm/llm-api.api +++ b/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 } diff --git a/llm-api/src/commonMain/kotlin/sk/ainet/llm/api/EmbeddingModel.kt b/llm-api/src/commonMain/kotlin/sk/ainet/llm/api/EmbeddingModel.kt index 625bb58b..910ccf90 100644 --- a/llm-api/src/commonMain/kotlin/sk/ainet/llm/api/EmbeddingModel.kt +++ b/llm-api/src/commonMain/kotlin/sk/ainet/llm/api/EmbeddingModel.kt @@ -20,6 +20,20 @@ public interface EmbeddingModel : AutoCloseable { public fun embed(texts: List): List = 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): List = embed(texts) + /** Output vector dimensionality. */ public val dimensions: Int diff --git a/llm-inference/bert/api/jvm/bert.api b/llm-inference/bert/api/jvm/bert.api index 57b3bad3..16fd09ae 100644 --- a/llm-inference/bert/api/jvm/bert.api +++ b/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 (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 (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 (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 (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; diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertConfig.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertConfig.kt index a130a16f..ac1aeb55 100644 --- a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertConfig.kt +++ b/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 + } } diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertEncoderRuntime.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertEncoderRuntime.kt index 1c4df175..dc345f2c 100644 --- a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertEncoderRuntime.kt +++ b/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( private val projectionWeight: Tensor? = null, private val projectionBias: Tensor? = 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( } /** - * 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( + 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( - 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 createBertEncoderRuntime( resolver: WeightNameResolver = BertSafeTensorsNameResolver(), mode: BertExecutionMode = BertExecutionMode.DIRECT, debug: Boolean = false, + pooling: BertPooling = BertPooling.MEAN, ): BertEncoderRuntime { val model = bertNetwork(config) @@ -337,5 +359,6 @@ public inline fun createBertEncoderRuntime( projectionWeight = projectionWeight, projectionBias = projectionBias, mode = mode, + pooling = pooling, ) } diff --git a/llm-inference/bert/src/commonTest/kotlin/sk/ainet/models/bert/BertPoolingConfigTest.kt b/llm-inference/bert/src/commonTest/kotlin/sk/ainet/models/bert/BertPoolingConfigTest.kt new file mode 100644 index 00000000..f133032c --- /dev/null +++ b/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 { + BertConfigParser.parsePooling("""{"pooling_mode_max_tokens": true}""") + } + assertFailsWith { + BertConfigParser.parsePooling("""{"pooling_mode_mean_sqrt_len_tokens": true}""") + } + } +} diff --git a/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertEncoderRuntimeTest.kt b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertEncoderRuntimeTest.kt index 2fa39153..d3097505 100644 --- a/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertEncoderRuntimeTest.kt +++ b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertEncoderRuntimeTest.kt @@ -187,4 +187,43 @@ class BertEncoderRuntimeTest { fun encode_emptyInputIsRejected() { assertFailsWith { runtime().encode(intArrayOf()) } } + + @Test + fun encode_clsPooling_takesFirstTokenRow() { + val cfg = config() + val rt = createBertEncoderRuntime(cfg, syntheticTensors(cfg), ctx, pooling = BertPooling.CLS) + val tokens = intArrayOf(1, 2, 3) + + // Expected: hidden-state row 0, L2-normalized (no projection configured). + val hidden = rt.forward(tokens) + val row = FloatArray(h) { c -> hidden.data[0, c] } + val norm = sqrt(row.sumOf { (it * it).toDouble() }).toFloat() + + val got = rt.encode(tokens) + assertEquals(h, got.size) + for (i in 0 until h) { + assertTrue(abs(row[i] / norm - got[i]) < 1e-5f, "CLS row mismatch at $i") + } + } + + @Test + fun encode_clsPooling_differsFromMeanOnMultiTokenInput() { + val cfg = config() + val cls = createBertEncoderRuntime(cfg, syntheticTensors(cfg), ctx, pooling = BertPooling.CLS) + .encode(intArrayOf(1, 2, 3)) + val mean = runtime().encode(intArrayOf(1, 2, 3)) + val diff = cls.indices.sumOf { abs(cls[it] - mean[it]).toDouble() } + assertTrue(diff > 1e-4, "CLS and MEAN pooling must differ on varied hidden states: diff=$diff") + } + + @Test + fun encode_clsPooling_ignoresAttentionMask() { + val cfg = config() + val rt = createBertEncoderRuntime(cfg, syntheticTensors(cfg), ctx, pooling = BertPooling.CLS) + val plain = rt.encode(intArrayOf(1, 2, 0)) + val masked = rt.encode(intArrayOf(1, 2, 0), attentionMask = intArrayOf(1, 1, 0)) + for (i in plain.indices) { + assertTrue(abs(plain[i] - masked[i]) < 1e-6f, "CLS pooling must not depend on the mask") + } + } } diff --git a/llm-providers/api/llm-providers.api b/llm-providers/api/llm-providers.api index 7d502d39..27077552 100644 --- a/llm-providers/api/llm-providers.api +++ b/llm-providers/api/llm-providers.api @@ -9,7 +9,42 @@ public final class sk/ainet/llm/providers/BertEmbeddingModel { public static final fun fromSafeTensors (Ljava/nio/file/Path;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/llm/api/EmbeddingModel; public static final fun fromSafeTensors (Ljava/nio/file/Path;Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/bert/BertExecutionMode;)Lsk/ainet/llm/api/EmbeddingModel; public static final fun fromSafeTensors (Ljava/nio/file/Path;Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/bert/BertExecutionMode;Ljava/lang/String;)Lsk/ainet/llm/api/EmbeddingModel; - public static synthetic fun fromSafeTensors$default (Ljava/nio/file/Path;Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/bert/BertExecutionMode;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/llm/api/EmbeddingModel; + public static final fun fromSafeTensors (Ljava/nio/file/Path;Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/bert/BertExecutionMode;Ljava/lang/String;Lsk/ainet/llm/providers/EmbeddingPrefixes;)Lsk/ainet/llm/api/EmbeddingModel; + public static synthetic fun fromSafeTensors$default (Ljava/nio/file/Path;Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/bert/BertExecutionMode;Ljava/lang/String;Lsk/ainet/llm/providers/EmbeddingPrefixes;ILjava/lang/Object;)Lsk/ainet/llm/api/EmbeddingModel; +} + +public final class sk/ainet/llm/providers/EmbeddingModelProfiles { + public static final field INSTANCE Lsk/ainet/llm/providers/EmbeddingModelProfiles; + public final fun apply (Lsk/ainet/llm/api/EmbeddingModel;Ljava/lang/String;)Lsk/ainet/llm/api/EmbeddingModel; + public final fun forRepo (Ljava/lang/String;)Lsk/ainet/llm/providers/EmbeddingPrefixes; +} + +public final class sk/ainet/llm/providers/EmbeddingPrefixes { + public static final field Companion Lsk/ainet/llm/providers/EmbeddingPrefixes$Companion; + public fun (Ljava/lang/String;Ljava/lang/String;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/String; + public final fun copy (Ljava/lang/String;Ljava/lang/String;)Lsk/ainet/llm/providers/EmbeddingPrefixes; + public static synthetic fun copy$default (Lsk/ainet/llm/providers/EmbeddingPrefixes;Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/llm/providers/EmbeddingPrefixes; + public fun equals (Ljava/lang/Object;)Z + public final fun getDocument ()Ljava/lang/String; + public final fun getQuery ()Ljava/lang/String; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/llm/providers/EmbeddingPrefixes$Companion { + public final fun getNONE ()Lsk/ainet/llm/providers/EmbeddingPrefixes; +} + +public final class sk/ainet/llm/providers/PrefixedEmbeddingModel : sk/ainet/llm/api/EmbeddingModel { + public fun (Lsk/ainet/llm/api/EmbeddingModel;Lsk/ainet/llm/providers/EmbeddingPrefixes;)V + public fun call (Lsk/ainet/llm/api/EmbeddingRequest;)Lsk/ainet/llm/api/EmbeddingResponse; + public fun close ()V + public fun embedDocument (Ljava/lang/String;)[F + public fun embedDocuments (Ljava/util/List;)Ljava/util/List; + public fun embedQuery (Ljava/lang/String;)[F + public fun getDimensions ()I } public final class sk/ainet/llm/providers/SkaiNetChatModel : sk/ainet/llm/api/StreamingChatModel { diff --git a/llm-providers/src/main/kotlin/sk/ainet/llm/providers/BertEmbeddingModel.kt b/llm-providers/src/main/kotlin/sk/ainet/llm/providers/BertEmbeddingModel.kt index a7c5dadc..d0f0a284 100644 --- a/llm-providers/src/main/kotlin/sk/ainet/llm/providers/BertEmbeddingModel.kt +++ b/llm-providers/src/main/kotlin/sk/ainet/llm/providers/BertEmbeddingModel.kt @@ -10,7 +10,6 @@ import sk.ainet.data.source.DataSourceException import sk.ainet.data.source.DataSourceRequest import sk.ainet.data.source.JvmDataSourceResolver import sk.ainet.io.JvmRandomAccessSource -import sk.ainet.io.safetensors.SafeTensorsParametersLoader import sk.ainet.llm.api.EmbeddingModel import sk.ainet.models.bert.BertConfigParser import sk.ainet.models.bert.BertExecutionMode @@ -59,6 +58,10 @@ public object BertEmbeddingModel { * @param mode eager [BertExecutionMode.DIRECT] (default) or traced/fused * [BertExecutionMode.OPTIMIZED] * @param modelId reported via [EmbeddingModel] response metadata + * @param prefixes query/document instruction prefixes for retrieval-tuned + * models; local snapshots carry no repo id to look up in + * [EmbeddingModelProfiles], so pass them explicitly (or use + * [fromHuggingFace], which applies the profile automatically) */ @JvmStatic @JvmOverloads @@ -67,6 +70,7 @@ public object BertEmbeddingModel { ctx: ExecutionContext = PooledExecutionContext(DirectCpuExecutionContext()), mode: BertExecutionMode = BertExecutionMode.DIRECT, modelId: String? = modelDir.fileName?.toString(), + prefixes: EmbeddingPrefixes = EmbeddingPrefixes.NONE, ): EmbeddingModel { val weightsFile = resolveWeightsFile(modelDir) val denseWeights = modelDir.resolve("2_Dense").resolve("model.safetensors") @@ -79,23 +83,32 @@ public object BertEmbeddingModel { if (denseConfig.exists()) denseConfig.readText() else null, ) + val poolingConfig = modelDir.resolve("1_Pooling").resolve("config.json") + val pooling = BertConfigParser.parsePooling( + if (poolingConfig.exists()) poolingConfig.readText() else null, + ) + val tokenizer = loadTokenizer(modelDir) + // FloatSafeTensorsLoader (not the engine's SafeTensorsParametersLoader): + // BGE-style checkpoints persist an I64 `embeddings.position_ids` buffer + // that the engine loader can't skip and our index-free encoder never needs. val loaders = buildList { - add(SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(weightsFile.toString()) })) + add(FloatSafeTensorsLoader(sourceProvider = { JvmRandomAccessSource.open(weightsFile.toString()) })) if (denseWeights.exists()) { - add(SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(denseWeights.toString()) })) + add(FloatSafeTensorsLoader(sourceProvider = { JvmRandomAccessSource.open(denseWeights.toString()) })) } } val tensors = runBlocking { BertNetworkLoader.loadWeightTensors(loaders, ctx, sk.ainet.lang.types.FP32::class) } - val runtime = createBertEncoderRuntime(config, tensors, ctx, mode = mode) + val runtime = createBertEncoderRuntime(config, tensors, ctx, mode = mode, pooling = pooling) - return SkaiNetEmbeddingModel( + val model: EmbeddingModel = SkaiNetEmbeddingModel( runtime = runtime, tokenizer = tokenizer, modelId = modelId, ) + return if (prefixes == EmbeddingPrefixes.NONE) model else PrefixedEmbeddingModel(model, prefixes) } /** @@ -116,7 +129,10 @@ public object BertEmbeddingModel { mode: BertExecutionMode = BertExecutionMode.DIRECT, ): EmbeddingModel { val modelDir = downloadSnapshot(repoId, revision) - return fromSafeTensors(modelDir, ctx, mode, modelId = repoId) + return fromSafeTensors( + modelDir, ctx, mode, modelId = repoId, + prefixes = EmbeddingModelProfiles.forRepo(repoId), + ) } // ------------------------------------------------------------------ @@ -128,6 +144,7 @@ public object BertEmbeddingModel { private val OPTIONAL_FILES = listOf( "vocab.txt", "tokenizer.json", + "1_Pooling/config.json", "2_Dense/config.json", "2_Dense/model.safetensors", ) diff --git a/llm-providers/src/main/kotlin/sk/ainet/llm/providers/FloatSafeTensorsLoader.kt b/llm-providers/src/main/kotlin/sk/ainet/llm/providers/FloatSafeTensorsLoader.kt new file mode 100644 index 00000000..e6346c45 --- /dev/null +++ b/llm-providers/src/main/kotlin/sk/ainet/llm/providers/FloatSafeTensorsLoader.kt @@ -0,0 +1,76 @@ +package sk.ainet.llm.providers + +import sk.ainet.context.ExecutionContext +import sk.ainet.io.ParametersLoader +import sk.ainet.io.RandomAccessSource +import sk.ainet.io.model.DataType +import sk.ainet.io.safetensors.StreamingSafeTensorsReader +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.DType +import java.nio.ByteBuffer +import java.nio.ByteOrder +import kotlin.reflect.KClass + +/** + * [ParametersLoader] over a SafeTensors file that loads *float weights only*, + * silently skipping integer buffers. + * + * Why this exists: HF BERT checkpoints exported by older `transformers` + * versions persist non-weight buffers — `embeddings.position_ids` (I64 + * `arange`) in BGE among others. The engine's `SafeTensorsParametersLoader` + * has no tensor filter and fails the whole load on the I64 → FP32 conversion. + * Our DSL `BertEmbeddings` is index-free, so these buffers are never needed. + * + * Interim: remove once the engine's loader accepts a tensor-name filter + * (SKaiNET-developers/SKaiNET#822) and route BERT loads back through + * `SafeTensorsParametersLoader`. + */ +internal class FloatSafeTensorsLoader( + private val sourceProvider: () -> RandomAccessSource, +) : ParametersLoader { + + @Suppress("UNCHECKED_CAST") + override suspend fun load( + ctx: ExecutionContext, + dtype: KClass, + onTensorLoaded: (name: String, tensor: Tensor) -> Unit, + ) { + StreamingSafeTensorsReader.open(sourceProvider()).use { reader -> + for (info in reader.tensors) { + val floats = when (info.dataType) { + DataType.FLOAT32 -> f32(reader.loadTensorData(info)) + DataType.FLOAT64 -> f64(reader.loadTensorData(info)) + DataType.FLOAT16 -> f16(reader.loadTensorData(info)) + DataType.BFLOAT16 -> bf16(reader.loadTensorData(info)) + // Non-float buffers (position_ids & friends) are never + // BERT weights — skip instead of failing the load. + else -> continue + } + val shape = Shape(*info.shape.map { it.toInt() }.toIntArray()) + val tensor = ctx.fromFloatArray(shape, dtype, floats) as Tensor + onTensorLoaded(info.name, tensor) + } + } + } + + private fun f32(bytes: ByteArray): FloatArray { + val buf = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).asFloatBuffer() + return FloatArray(buf.remaining()).also { buf.get(it) } + } + + private fun f64(bytes: ByteArray): FloatArray { + val buf = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).asDoubleBuffer() + return FloatArray(buf.remaining()) { buf.get(it).toFloat() } + } + + private fun f16(bytes: ByteArray): FloatArray { + val buf = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).asShortBuffer() + return FloatArray(buf.remaining()) { java.lang.Float.float16ToFloat(buf.get(it)) } + } + + private fun bf16(bytes: ByteArray): FloatArray { + val buf = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).asShortBuffer() + return FloatArray(buf.remaining()) { Float.fromBits(buf.get(it).toInt() shl 16) } + } +} diff --git a/llm-providers/src/main/kotlin/sk/ainet/llm/providers/PrefixedEmbeddingModel.kt b/llm-providers/src/main/kotlin/sk/ainet/llm/providers/PrefixedEmbeddingModel.kt new file mode 100644 index 00000000..1a4e155a --- /dev/null +++ b/llm-providers/src/main/kotlin/sk/ainet/llm/providers/PrefixedEmbeddingModel.kt @@ -0,0 +1,89 @@ +package sk.ainet.llm.providers + +import sk.ainet.llm.api.EmbeddingModel +import sk.ainet.llm.api.EmbeddingRequest +import sk.ainet.llm.api.EmbeddingResponse + +/** + * Per-role instruction prefixes for retrieval-tuned embedding models — the + * asymmetry between "embed this query" and "embed this passage" that models + * like E5 *require* and BGE recommend for short queries. + */ +public data class EmbeddingPrefixes( + val query: String, + val document: String, +) { + public companion object { + /** No prefixes — symmetric models (LEAF, MiniLM). */ + public val NONE: EmbeddingPrefixes = EmbeddingPrefixes("", "") + } +} + +/** + * Decorates an [EmbeddingModel] with per-role prefixes: [embedQuery] prepends + * [EmbeddingPrefixes.query], [embedDocument]/[embedDocuments] prepend + * [EmbeddingPrefixes.document]. Untyped calls ([call], [embed]) use the + * *document* role — indexing is the bulk path, and an unprefixed E5 embedding + * is wrong for both roles, so defaulting to document keeps `embed(chunks)` + * correct. + */ +public class PrefixedEmbeddingModel( + private val delegate: EmbeddingModel, + private val prefixes: EmbeddingPrefixes, +) : EmbeddingModel { + + override fun call(request: EmbeddingRequest): EmbeddingResponse = + delegate.call(request.copy(inputs = request.inputs.map { prefixes.document + it })) + + override fun embedQuery(text: String): FloatArray = + delegate.embed(prefixes.query + text) + + override fun embedDocument(text: String): FloatArray = + delegate.embed(prefixes.document + text) + + override fun embedDocuments(texts: List): List = + delegate.embed(texts.map { prefixes.document + it }) + + override val dimensions: Int get() = delegate.dimensions + + override fun close(): Unit = delegate.close() +} + +/** + * Known per-model prefix profiles, keyed by Hugging Face repo-id prefix + * (case-insensitive). [BertEmbeddingModel.fromHuggingFace] consults this + * registry automatically; local-snapshot loads pass prefixes explicitly. + */ +public object EmbeddingModelProfiles { + + private val profiles: List> = listOf( + // E5: both roles REQUIRE prefixes (intfloat model card). + "intfloat/multilingual-e5" to EmbeddingPrefixes(query = "query: ", document = "passage: "), + "intfloat/e5" to EmbeddingPrefixes(query = "query: ", document = "passage: "), + // BGE v1.5 English: query instruction recommended for short queries; passages plain. + "BAAI/bge-small-en" to EmbeddingPrefixes( + query = "Represent this sentence for searching relevant passages: ", + document = "", + ), + "BAAI/bge-base-en" to EmbeddingPrefixes( + query = "Represent this sentence for searching relevant passages: ", + document = "", + ), + "BAAI/bge-large-en" to EmbeddingPrefixes( + query = "Represent this sentence for searching relevant passages: ", + document = "", + ), + ) + + /** Prefixes for [repoId], or [EmbeddingPrefixes.NONE] when unknown. */ + public fun forRepo(repoId: String): EmbeddingPrefixes = + profiles.firstOrNull { (prefix, _) -> repoId.startsWith(prefix, ignoreCase = true) } + ?.second + ?: EmbeddingPrefixes.NONE + + /** Wrap [model] with the profile for [repoId]; returns [model] unchanged when none applies. */ + public fun apply(model: EmbeddingModel, repoId: String): EmbeddingModel { + val prefixes = forRepo(repoId) + return if (prefixes == EmbeddingPrefixes.NONE) model else PrefixedEmbeddingModel(model, prefixes) + } +} diff --git a/llm-providers/src/test/kotlin/sk/ainet/llm/providers/PrefixedEmbeddingModelTest.kt b/llm-providers/src/test/kotlin/sk/ainet/llm/providers/PrefixedEmbeddingModelTest.kt new file mode 100644 index 00000000..d059a1c3 --- /dev/null +++ b/llm-providers/src/test/kotlin/sk/ainet/llm/providers/PrefixedEmbeddingModelTest.kt @@ -0,0 +1,67 @@ +package sk.ainet.llm.providers + +import sk.ainet.llm.api.Embedding +import sk.ainet.llm.api.EmbeddingModel +import sk.ainet.llm.api.EmbeddingRequest +import sk.ainet.llm.api.EmbeddingResponse +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertSame + +class PrefixedEmbeddingModelTest { + + /** Records every text handed to the delegate. */ + private class RecordingModel : EmbeddingModel { + val seen = mutableListOf() + override fun call(request: EmbeddingRequest): EmbeddingResponse { + seen += request.inputs + return EmbeddingResponse( + request.inputs.mapIndexed { i, _ -> Embedding(i, floatArrayOf(1f)) }, + ) + } + override val dimensions: Int = 1 + } + + private val e5 = EmbeddingPrefixes(query = "query: ", document = "passage: ") + + @Test + fun embedQuery_prependsQueryPrefix() { + val delegate = RecordingModel() + PrefixedEmbeddingModel(delegate, e5).embedQuery("what is a pangram?") + assertEquals(listOf("query: what is a pangram?"), delegate.seen) + } + + @Test + fun embedDocument_andBatch_prependDocumentPrefix() { + val delegate = RecordingModel() + val model = PrefixedEmbeddingModel(delegate, e5) + model.embedDocument("a") + model.embedDocuments(listOf("b", "c")) + assertEquals(listOf("passage: a", "passage: b", "passage: c"), delegate.seen) + } + + @Test + fun untypedCalls_useDocumentRole() { + val delegate = RecordingModel() + val model = PrefixedEmbeddingModel(delegate, e5) + model.embed("a") + model.call(EmbeddingRequest(listOf("b"))) + assertEquals(listOf("passage: a", "passage: b"), delegate.seen) + } + + @Test + fun profiles_e5AndBgeResolve_caseInsensitive() { + assertEquals("query: ", EmbeddingModelProfiles.forRepo("intfloat/multilingual-e5-small").query) + assertEquals("passage: ", EmbeddingModelProfiles.forRepo("INTFLOAT/Multilingual-E5-Large").document) + val bge = EmbeddingModelProfiles.forRepo("BAAI/bge-small-en-v1.5") + assertEquals("Represent this sentence for searching relevant passages: ", bge.query) + assertEquals("", bge.document) + } + + @Test + fun profiles_unknownRepo_returnsModelUnwrapped() { + assertEquals(EmbeddingPrefixes.NONE, EmbeddingModelProfiles.forRepo("MongoDB/mdbr-leaf-ir")) + val delegate = RecordingModel() + assertSame(delegate, EmbeddingModelProfiles.apply(delegate, "MongoDB/mdbr-leaf-ir")) + } +}