BERT on the DSL path: complete bertNetwork(), BertEncoderRuntime (eager + traced), one-call BertEmbeddingModel, legacy runtime removed

CHANGELOG.md

@@ -7,6 +7,58 @@ version line is kept in lock-step with the underlying SKaiNET engine
 The format roughly follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+
+- **BERT sentence embeddings completed on the DSL path.** `bertNetwork()` is now a numerically
+  complete `tokens → hidden-states` encoder: the new `BertEmbeddings` module adds absolute-position
+  and token-type embeddings (index-free `narrow`-based lookups, single-segment) that the DSL
+  definition previously omitted. New `BertEncoderRuntime` executes it eagerly (**DIRECT**, default)
+  or as a traced, optimized **ComputeGraph** (**OPTIMIZED**, shape-specialized per sequence length
+  with an LRU cache) and adds masked mean pooling, the optional sentence-transformers `2_Dense`
+  projection, and L2 normalization on top of the pure encoder graph. The encoder trace lowers to
+  StableHLO (gather / dot_general / SDPA preserved) — export is gate-tested; IREE *execution* of the
+  exported module stays out of scope for now. Verified against the PyTorch-validated legacy runtime
+  on real MongoDB/mdbr-leaf-mt (hidden-state parity ≤ 2.2e-6) and DIRECT-vs-OPTIMIZED bit-exact.
+- **One-call embedding factory with built-in Hugging Face download.**
+  `BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-mt")` (llm-providers) downloads the
+  snapshot via the engine's `skainet-data-source` (`hf://` URIs, `HF_TOKEN`-aware) into
+  `~/.cache/skainet/models/`, streamed with `.part` + atomic rename, offline-safe after the first
+  run; `fromSafeTensors(dir)` loads a local snapshot, auto-detecting weights, config, tokenizer
+  (`vocab.txt` → `tokenizer.json`), and the `2_Dense/` head. `kbert-cli` accepts an HF repo id
+  directly: `kbert MongoDB/mdbr-leaf-mt "query" "doc"`.
+- `BertConfigParser` — shared `config.json` (+ `2_Dense/config.json` → `projectionDim`) parser,
+  consolidating the copies previously living in `KBertJava` and downstream apps.
+
+### Fixed
+
+- **BERT post-norm residual wiring.** The single-block-per-layer `bertNetwork()` definition wired the
+  FFN residual to the pre-LayerNorm value — the transformer blocks' residual rule fits pre-norm
+  decoder stacks, but BERT is post-norm. Each encoder layer is now two blocks (`attn` / `ffn`) so
+  every residual segment starts at the correct value.
+- **Bias-free `2_Dense` projection heads were silently dropped.** The legacy eager runtime required
+  projection weight *and* bias; LEAF models ship `bias=false`, so it skipped the projection entirely
+  (returning 384-dim vectors while advertising 1024). `BertEncoderRuntime` applies bias-free
+  projections; `KBertJava` now picks up `2_Dense/` heads it previously ignored.
+- **Graph replay dropped `permute` axes.** The ComputeGraph executor's builtin dispatch replayed
+  `permute` as a plain last-two-dims transpose; `LLMFusedOpHandlers` registers an axes-aware
+  `permute` handler (the registry precedes the builtin), fixing every multi-token attention trace —
+  single-token decode never hit it. Remove once the engine executor honors `axes` upstream.
+
+### Removed
+
+- **BREAKING: the deprecated hand-coded BERT stack is gone** — `BertRuntime`, `BertRuntimeWeights`,
+  `BertLayerWeights`, `loadBertWeights`, `BertWeightMapper`, `BertTensorNames`, `BertIngestion`, and
+  `BertNetworkLoader.fromRuntimeWeights`. Migrate to `createBertEncoderRuntime(config, tensors, ctx)`
+  (tensors from `BertNetworkLoader.loadWeightTensors`) or, one level up, to
+  `BertEmbeddingModel.fromSafeTensors(...)` / `fromHuggingFace(...)`. `SkaiNetEmbeddingModel`'s
+  constructor now takes `BertEncoderRuntime`; `KBertJava` / `KBertSession` keep their method surface
+  (`loadSafeTensors` / `encode` / `similarity`) with the constructor type changing. `BertModelConfig`
+  and `MDBR_LEAF_IR_CONFIG` moved to `BertConfig.kt` (same package — imports unaffected). The
+  `docs/optimizable-LLM-NNs-DAG.md` reference in the old deprecation pointed at a document that never
+  existed; the real migration guide is `explanation/dsl-vs-handcoded.adoc`.
+
 ## [0.35.0] — 2026-07-09
 
 Ships against **SKaiNET engine 0.35.0**, whose new `argMax` op this release uses to fold the LLM

README.md

@@ -88,7 +88,7 @@ Honest status — see the project-status note at the top of this README.
 | **Qwen 2 / 3** | DSL + loaders present; runs through the shared decoder path. Early; Qwen3 RoPE / QK-norm fixes landed in 0.23.2. |
 | **Gemma 2 / 3 / 3n** | DSL + loaders present (Gemma 4 via the SafeTensors path); has the most test coverage, but not verified end-to-end. |
 | **Apertus** | DSL + loaders present; declared end-to-end in 0.23.1, still early. |
-| **BERT** | Encoder for embeddings only — no text generation, no tool calling. |
+| **BERT** | Sentence embeddings on the DSL path (`bertNetwork()` + `BertEncoderRuntime`, eager or traced/fused) — verified against sentence-transformers on MongoDB/mdbr-leaf. One-call `BertEmbeddingModel.fromHuggingFace(...)` with built-in Hub download. No text generation, no tool calling. |
 | **Voxtral** | TTS / voice; architecture code only — no runtime facade or CLI yet. |
 
 ### Near term

docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc

@@ -108,7 +108,7 @@ The goal is to extend the DSL to support these patterns over time.
 
 |BERT
 |`bertNetwork()`
-|`BertRuntime` deprecated
+|DSL-only: `BertEncoderRuntime` (eager + traced graph); legacy runtime removed
 
 |Voxtral
 |`voxtralBackboneNetwork()`

docs/modules/ROOT/pages/explanation/embeddings.adoc

@@ -1,7 +1,7 @@
 = How Embeddings Work
 :description: What the BERT runtime actually computes, why mean pooling + L2 normalization, and where the dimensionality comes from.
 
-This page is the "why" companion to xref:tutorials/embeddings.adoc[Embeddings — Getting Started]. It explains what a `BertRuntime.forward(...)` actually returns, why the public API performs mean pooling and L2 normalization on your behalf, and why two embeddings of the same dimension can still be incompatible.
+This page is the "why" companion to xref:tutorials/embeddings.adoc[Embeddings — Getting Started]. It explains what a `BertEncoderRuntime.forward(...)` actually returns, why the public API performs mean pooling and L2 normalization on your behalf, and why two embeddings of the same dimension can still be incompatible.
 
 == From Tokens to a Single Vector
 
@@ -12,7 +12,7 @@ For most downstream uses (semantic search, similarity, RAG retrieval) you want a
 * *`[CLS]` pooling*: take the embedding of the first token (the special `[CLS]` token BERT prepends). Original BERT was trained to use this for classification.
 * *Mean pooling*: average across the token dimension. Sentence-transformers (and `all-MiniLM-L6-v2`, `all-mpnet-base-v2`, etc.) are *fine-tuned* for this — they outperform `[CLS]` for similarity tasks, often by a wide margin.
 
-The `BertRuntime` implements *mean pooling* because that's what every modern sentence-encoder GGUF in the wild expects. The `EmbeddingModel.embed(...)` methods return the pooled vector directly.
+The `BertEncoderRuntime` implements *mean pooling* because that's what every modern sentence-encoder checkpoint in the wild expects. The `EmbeddingModel.embed(...)` methods return the pooled vector directly.
 
 [mermaid]
 ....
@@ -33,7 +33,7 @@ After mean pooling you get a vector whose *magnitude* depends on the input — l
 * if both vectors are L2-normalized to length 1, that division is 1, and cosine similarity reduces to a plain dot product;
 * most vector stores assume normalized inputs and compute cosine as dot product internally — passing un-normalized vectors silently gives the wrong rankings.
 
-`BertRuntime` therefore L2-normalizes the pooled vector before returning it. The Java/Kotlin `embed(...)` calls always return unit-length vectors.
+`BertEncoderRuntime` therefore L2-normalizes the pooled vector (after the optional sentence-transformers `2_Dense` projection) before returning it. The Java/Kotlin `embed(...)` calls always return unit-length vectors.
 
 == Where Does the Dimension Come From?
 
@@ -72,16 +72,22 @@ For practical retrieval, *thresholds vary by encoder* — a 0.7 cutoff on `all-M
 
 == Tokenization Matters
 
-The same text can produce different vectors with different tokenizers — this is not a bug, it's the model's input representation changing. Always pair an encoder with the tokenizer it was trained with. `KBertJava.loadGGUF(...)` reads the tokenizer config out of the GGUF metadata so you can't accidentally mix them up; if you wire `BertRuntime` and `HuggingFaceTokenizer` manually in Kotlin, make sure both come from the same model.
+The same text can produce different vectors with different tokenizers — this is not a bug, it's the model's input representation changing. Always pair an encoder with the tokenizer it was trained with. `BertEmbeddingModel.fromHuggingFace(...)` / `fromSafeTensors(...)` load the tokenizer from the same snapshot as the weights so you can't accidentally mix them up; if you wire `BertEncoderRuntime` and `HuggingFaceTokenizer` manually in Kotlin, make sure both come from the same model.
 
 == Where the Code Lives
 
 [cols="2,3"]
 |===
 | Concern | Source
 
-| Encoder forward pass + mean pooling + L2 norm
-| `llm-inference/bert/.../BertRuntime.kt`
+| Network definition (DSL)
+| `llm-inference/bert/.../BertNetworkDef.kt` + `BertEmbeddings.kt`
+
+| Encoder runtime: forward + mean pooling + projection + L2 norm
+| `llm-inference/bert/.../BertEncoderRuntime.kt`
+
+| One-call factory (local + Hugging Face download)
+| `llm-providers/.../BertEmbeddingModel.kt`
 
 | Provider-neutral SPI
 | `llm-api/.../EmbeddingModel.kt`

docs/modules/ROOT/pages/tutorials/embeddings.adoc

@@ -6,7 +6,7 @@ This tutorial walks through producing dense vector embeddings for text — the k
 == Prerequisites
 
 * JDK 21+ (Java 25 preferred for the Vector API)
-* A BERT-family GGUF model on disk, e.g. `all-MiniLM-L6-v2.gguf` (384 dims, ~80 MB) or any sentence-transformers model converted to GGUF
+* Nothing else — models download from the Hugging Face Hub on first use and are cached under `~/.cache/skainet/models/` (or point at a local sentence-transformers SafeTensors snapshot)
 
 == From the CLI
 
@@ -15,15 +15,15 @@ The fastest way to verify embeddings work end-to-end:
 [source,bash]
 ----
 ./gradlew :llm-apps:kbert-cli:run \
-  --args="all-MiniLM-L6-v2.gguf 'The quick brown fox jumps over the lazy dog'"
+  --args="MongoDB/mdbr-leaf-mt 'The quick brown fox jumps over the lazy dog'"
 ----
 
-Or with a document for similarity:
+Or with a document for similarity (a local snapshot directory works in place of the repo id):
 
 [source,bash]
 ----
 ./gradlew :llm-apps:kbert-cli:run \
-  --args="all-MiniLM-L6-v2.gguf 'pangram' 'A pangram is a sentence that contains every letter of the alphabet.'"
+  --args="MongoDB/mdbr-leaf-mt 'pangram' 'A pangram is a sentence that contains every letter of the alphabet.'"
 ----
 
 == From Kotlin / Java — `EmbeddingModel` SPI
@@ -40,24 +40,20 @@ public interface EmbeddingModel : AutoCloseable {
 }
 ----
 
-Adapter wiring `BertRuntime` to the SPI lives in `llm-providers/SkaiNetEmbeddingModel.kt`:
+The one-call factory in `llm-providers/BertEmbeddingModel.kt` builds the whole stack — DSL network (`bertNetwork()`), weight mapping, tokenizer, encoder runtime — behind the SPI:
 
 [source,kotlin]
 ----
-import sk.ainet.llm.providers.SkaiNetEmbeddingModel
-import sk.ainet.models.bert.BertIngestion
+import sk.ainet.llm.providers.BertEmbeddingModel
 
-val ingestion = BertIngestion<FP32>(ctx = ctx, dtype = FP32::class)
-val weights = ingestion.loadGguf { JvmRandomAccessSource.open("all-MiniLM-L6-v2.gguf") }
-val runtime = BertRuntime(ctx, weights, FP32::class)
-val tokenizer = HuggingFaceTokenizer.fromGguf(weights.metadata)
+// Straight from the Hugging Face Hub (downloads + caches on first use;
+// picks up HF_TOKEN for gated repos):
+val model: EmbeddingModel = BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-mt")
 
-val model: EmbeddingModel = SkaiNetEmbeddingModel(
-    runtime = runtime,
-    tokenizer = tokenizer,
-    dimensions = weights.metadata.embeddingLength,
-    modelId = "all-MiniLM-L6-v2",
-)
+// Or from a local sentence-transformers snapshot directory
+// (auto-detects config.json, vocab.txt / tokenizer.json, model.safetensors,
+// and the optional 2_Dense/ projection head):
+val local: EmbeddingModel = BertEmbeddingModel.fromSafeTensors(Path.of("/models/mdbr-leaf-mt"))
 
 // Single text — convenience overload.
 val vector: FloatArray = model.embed("The quick brown fox")
@@ -84,16 +80,23 @@ fun cosine(a: FloatArray, b: FloatArray): Float {
 
 == From Java
 
-`KBertJava` exposes the same surface for pure-Java consumers:
+`BertEmbeddingModel` is `@JvmStatic` throughout, and `KBertJava` offers a smaller session-style surface for pure-Java consumers:
 
 [source,java]
 ----
+import sk.ainet.llm.api.EmbeddingModel;
+import sk.ainet.llm.providers.BertEmbeddingModel;
+
+EmbeddingModel model = BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-mt");
+float[] vector = model.embed("The quick brown fox");
+
+// Or the session facade over a local snapshot:
 import sk.ainet.models.bert.java.KBertJava;
 import sk.ainet.models.bert.java.KBertSession;
 
-try (KBertSession session = KBertJava.loadGGUF(Path.of("all-MiniLM-L6-v2.gguf"))) {
-    float[] vector = session.embed("The quick brown fox");
-    System.out.println("dim=" + vector.length);
+try (KBertSession session = KBertJava.loadSafeTensors(Path.of("/models/mdbr-leaf-mt"))) {
+    float[] v = session.encode("The quick brown fox");
+    float sim = session.similarity("query text", "document text");
 }
 ----
 

docs/specs/spring-ai-adapter.md

@@ -39,7 +39,7 @@ On branch `feature/llm-api-neutral-spi`:
   `Flow` for streaming. Deps: `kotlin-stdlib` + `kotlinx-coroutines` only.
 - `llm-providers/` — JVM module with `SkaiNetChatModel<T>` (wraps any
   `InferenceRuntime<T>` + `Tokenizer` + `ChatTemplate`) and `SkaiNetEmbeddingModel<T>`
-  (wraps `BertRuntime<T>`).
+  (wraps `BertEncoderRuntime<T>`).
 - BOM updated; binary-compat baseline (`api/`) generated; basic unit tests for
   mappers and stop-sequence helper.
 
@@ -115,8 +115,8 @@ Conditional bean wiring:
   `ChatTemplate` via `ModelRegistry.detect(...)` (or the explicit override),
   wraps in `SkaiNetChatModel`, then in `SpringSkaiNetChatModel`.
 - `@Bean @ConditionalOnMissingBean EmbeddingModel skaiNetEmbeddingModel(...)` —
-  same pattern with `BertRuntime` (or its eventual `OptimizedLLMRuntime`-based
-  replacement) + `SkaiNetEmbeddingModel`.
+  same pattern with `BertEncoderRuntime` (the DSL-path encoder) +
+  `SkaiNetEmbeddingModel`, or simply `BertEmbeddingModel.fromHuggingFace(...)`.
 
 ## Open dependencies / blockers
 

gradle/libs.versions.toml

@@ -99,6 +99,8 @@ skainet-io-core = { module = "sk.ainet.core:skainet-io-core" }
 skainet-io-gguf = { module = "sk.ainet.core:skainet-io-gguf" }
 skainet-io-iree-params = { module = "sk.ainet.core:skainet-io-iree-params" }
 skainet-io-safetensors = { module = "sk.ainet.core:skainet-io-safetensors" }
+# URI-backed data sources (hf:// downloads) — BertEmbeddingModel.fromHuggingFace
+skainet-data-source = { module = "sk.ainet.core:skainet-data-source" }
 
 [plugins]
 androidLibrary = { id = "com.android.library", version.ref = "agp" }

llm-apps/kbert-cli/build.gradle.kts

@@ -12,6 +12,7 @@ dependencies {
     implementation(platform(project(":llm-bom")))
 
     implementation(project(":llm-inference:bert"))
+    implementation(project(":llm-providers"))
     implementation(libs.skainet.lang.core)
     implementation(libs.skainet.io.core)
     implementation(libs.skainet.io.safetensors)