diff --git a/CHANGELOG.md b/CHANGELOG.md index fdcb5c3f..dd025d3f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.36.0] — 2026-07-12 + +Ships against **SKaiNET engine 0.36.0**. Headline: **BERT is now completely defined on the DSL +path** — the legacy hand-coded eager stack is removed (**BREAKING**, see *Removed*), and sentence +embeddings get a one-call factory with built-in Hugging Face Hub download. Also new: a **T5 +encoder-decoder** runtime and a **vec2text embedding-inversion** pipeline (invert GTR embeddings +back to text). Downstream impact: +indexing the leaf-cli reference corpus (56 chunks) drops from 676.9 s to 44.5 s (~15×) with +identical embeddings. + ### Added - **BERT sentence embeddings completed on the DSL path.** `bertNetwork()` is now a numerically @@ -30,6 +40,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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. +- **T5 encoder-decoder runtime** (`llm-inference/t5`, `sk.ainet.models.t5`). Hand-coded in the + direct tensor-ops style (per-head attention via narrow/matmul/softmax, batch 1, no KV cache — + the greedy decoder recomputes the stack per step), handling T5's specifics: no 1/√d attention + scaling, learned relative-position bias (`T5RelativeBias`, block-0 table shared per stack, + none in cross-attention), RMSNorm-style T5LayerNorm, un-gated ReLU FFN, tied embeddings with + `d_model^-0.5` logit scaling. Includes `GtrEmbedder` — GTR sentence embeddings exactly as + vec2text consumes them (raw T5 encoder + mean pooling; deliberately no Dense projection and no + L2 normalization) — with a parity test against real `sentence-transformers/gtr-t5-base` weights. +- **vec2text embedding inversion** (`llm-inference/vec2text`, `sk.ainet.models.vec2text`). + Port of vec2text's greedy corrector loop (`sequence_beam_width = 1`): `InversionModel` + produces an initial hypothesis from a target GTR embedding, then `CorrectorModel` iteratively + re-embeds and corrects it, early-stopping when the cosine score plateaus — `Vec2TextInverter` + returns the best reconstruction plus the full step trace. Verified with an end-to-end + round-trip test on real gtr-base weights. ### Fixed @@ -42,9 +66,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 (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. + `permute` as a plain last-two-dims transpose, breaking every multi-token attention trace — + single-token decode never hit it. Fixed upstream in engine 0.36.0 + ([SKaiNET#803](https://github.com/SKaiNET-developers/SKaiNET/pull/803)), which this release + consumes; the interim axes-aware `permute` handler in `LLMFusedOpHandlers` (never in a published + release) is removed again. ### Removed @@ -742,6 +768,7 @@ Version-aligned with **SKaiNET 0.21.0**. Last published transformers release before the engine-aligned version line. See `git log v0.16.0..0.18.0` for details. +[0.36.0]: https://github.com/SKaiNET-developers/SKaiNET-transformers/releases/tag/0.36.0 [0.31.0]: https://github.com/SKaiNET-developers/SKaiNET-transformers/releases/tag/0.31.0 [0.30.0]: https://github.com/SKaiNET-developers/SKaiNET-transformers/releases/tag/0.30.0 [0.28.1]: https://github.com/SKaiNET-developers/SKaiNET-transformers/releases/tag/0.28.1 diff --git a/README.md b/README.md index 5d4c6291..fecd08c4 100644 --- a/README.md +++ b/README.md @@ -74,9 +74,11 @@ flowchart LR HLO --> Native["Native code"] ``` -Today every model family runs through the **eager JVM path**. The StableHLO / -native path is shared with the engine and not yet wired for full transformer -models. +The **eager JVM path** is the primary way every model family runs today. The +StableHLO / native path is shared with the engine and wired for the first +families: FunctionGemma exports a compiled edge build (0.35.0), and the BERT +encoder traces to an optimized ComputeGraph and lowers to StableHLO (0.36.0); +full generative-model coverage is still in progress. ### Where each architecture fits @@ -89,6 +91,7 @@ Honest status — see the project-status note at the top of this README. | **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** | 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. | +| **T5 / GTR** | Encoder-decoder runtime (hand-coded, batch 1, no KV cache) + `GtrEmbedder`, powering the **vec2text** embedding-inversion pipeline — verified with a real-weights gtr-base round-trip test. | | **Voxtral** | TTS / voice; architecture code only — no runtime facade or CLI yet. | ### Near term @@ -103,15 +106,38 @@ Honest status — see the project-status note at the top of this README. ## Current release -The current release is **0.35.0** (against **SKaiNET 0.35.0**) — it adds **FunctionGemma** -self-compiled from the SKaiNET NN DSL: a one-dependency function-calling sLLM -(`skainet-transformers-runtime-kgemma`) with an eager one-line API -(`FunctionGemma.fromGguf(gguf).call("turn the light on")` → `ToolCall(set_lights, {state="on"})`, runs -anywhere on CPU, no iree) **and** a no-Python compiled edge export (`FunctionGemma.exportCompiled` / -`compile-gemma.sh`) verified token-for-token against llama.cpp on the SL2610 board. It uses the engine's -new `argMax` op to fold the `logits → token-ids` argmax tail into the DSL trace. - -It builds on **0.34.1** — a patch that layer-qualifies the +The current release is **0.36.0** (against **SKaiNET 0.36.0**) — **BERT is now completely +defined in the SKaiNET NN DSL**, and the deprecated hand-coded eager BERT stack is **removed +(BREAKING)** in the same release: + +- `bertNetwork()` is a numerically complete `tokens → hidden-states` encoder: the new + `BertEmbeddings` module adds the absolute-position and token-type embeddings the DSL definition + previously omitted, and each encoder layer is wired as two post-norm blocks so every residual + lands on the right value. +- `BertEncoderRuntime` runs the same definition **eagerly** (`DIRECT`, default) or as a traced, + LLM-pipeline-**optimized ComputeGraph** (`OPTIMIZED`, bit-exact vs eager), adds masked mean + pooling, the optional sentence-transformers `2_Dense` projection, and L2 normalization — and + `exportTape(...)` lowers the encoder to StableHLO. +- One-call consumption: `BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-mt")` / + `fromSafeTensors(dir)` behind the neutral `EmbeddingModel` SPI, with built-in Hub download + (`HF_TOKEN`-aware, cached, offline-safe after the first run). +- Downstream effect: indexing the leaf-cli reference corpus dropped **676.9 s → 44.5 s (~15×)** + with identical embeddings. Migration notes for the removed `BertRuntime` stack are in the + [CHANGELOG](CHANGELOG.md) and the + [BERT-as-DSL explanation](docs/modules/ROOT/pages/explanation/bert-dsl.adoc). + +0.36.0 also adds a **T5 encoder-decoder** runtime (`llm-inference/t5`) with `GtrEmbedder`, and a +**vec2text embedding-inversion** pipeline (`llm-inference/vec2text`) that iteratively reconstructs +text from a GTR embedding — verified end-to-end against real +`sentence-transformers/gtr-t5-base` weights. + +It builds on **0.35.0**, which added **FunctionGemma** self-compiled from the SKaiNET NN DSL: a +one-dependency function-calling sLLM (`skainet-transformers-runtime-kgemma`) with an eager one-line +API (`FunctionGemma.fromGguf(gguf).call("turn the light on")` → `ToolCall(set_lights, {state="on"})`, +runs anywhere on CPU, no iree) **and** a no-Python compiled edge export +(`FunctionGemma.exportCompiled` / `compile-gemma.sh`) verified token-for-token against llama.cpp on +the SL2610 board, using the engine's new `argMax` op to fold the `logits → token-ids` argmax tail +into the DSL trace; and on **0.34.1** — a patch that layer-qualifies the Moonshine encoder's attention/LayerNorm parameter names so by-name weight loading can tell the layers apart (no public API change) — and on **0.34.0**, which adds the first **Moonshine** speech-to-text encoder authored entirely in the SKaiNET NN DSL (`skainet-transformers-inference-moonshine`, @@ -138,7 +164,7 @@ The recommended way to consume is via the BOM. It pins every published `skainet- ```kotlin dependencies { - implementation(platform("sk.ainet.transformers:skainet-transformers-bom:0.35.0")) + implementation(platform("sk.ainet.transformers:skainet-transformers-bom:0.36.0")) // Versions resolved from the BOM: implementation("sk.ainet.transformers:skainet-transformers-core") @@ -165,7 +191,7 @@ dependencies { | `llm-api` | Framework-neutral interfaces (`ChatModel`, `EmbeddingModel`, `ToolDefinition`) — Spring AI-shaped. | | `transformer-core` | Framework NN primitives (attention, KV-cache family, embedding, norms, RoPE, FFNs, linear projection). `lang-core`-only → **all targets incl. `androidNative`**; re-exported by `llm-core`. | | `llm-core` | `OptimizedLLMRuntime`, `ModelRegistry`, `UnifiedModelLoader`, shared abstractions. | -| `llm-inference/` | Per-architecture network DSLs and weight loaders (`llama`, `gemma`, `qwen`, `apertus`, `bert`). | +| `llm-inference/` | Per-architecture network DSLs and weight loaders (`llama`, `gemma`, `qwen`, `apertus`, `bert`, `t5`, `vec2text`). | | `llm-runtime/` | Per-architecture runtime facades (`kllama`, `kgemma`, `kqwen`, `kapertus`). | | `llm-agent` | Chat templates, tool-call parsers, agent loops; Java surface. | | `llm-apps` | CLIs: `skainet-cli` (unified), `kllama-cli`, `kbert-cli`, plus `kllama-java-sample`. | @@ -195,6 +221,22 @@ java -jar skainet-all.jar -m model.gguf --agent --template=apertus `--template` accepts `llama3`, `chatml`, `qwen`, `gemma`, `apertus` (auto-detected from GGUF metadata if omitted). +### Embeddings: LEAF in one call + +Sentence embeddings with MongoDB's compact LEAF retrieval models need a single factory call — +the model downloads from the Hugging Face Hub and is cached on first use: + +```kotlin +import sk.ainet.llm.providers.BertEmbeddingModel + +BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-ir").use { model -> + val vector = model.embed("The quick brown fox") // L2-normalized FloatArray +} +``` + +See the [Getting Started with LEAF tutorial](docs/modules/ROOT/pages/tutorials/getting-started-leaf.adoc) +and the [BERT-as-DSL explanation](docs/modules/ROOT/pages/explanation/bert-dsl.adoc). + ### Java consumers ```java diff --git a/docs/modules/ROOT/nav.adoc b/docs/modules/ROOT/nav.adoc index e7b6c2c5..668739c6 100644 --- a/docs/modules/ROOT/nav.adoc +++ b/docs/modules/ROOT/nav.adoc @@ -6,6 +6,7 @@ * xref:tutorials/tool-calling.adoc[Tool Calling with Any Model] * xref:tutorials/llama3-tool-calling.adoc[Llama 3 / 3.1 / 3.2 Tool Calling] * xref:tutorials/embeddings.adoc[Embeddings — Getting Started] +* xref:tutorials/getting-started-leaf.adoc[Getting Started with LEAF] * xref:tutorials/smoke-tests.adoc[Running Smoke Tests] .How-to Guides @@ -26,6 +27,7 @@ .Explanation * xref:explanation/pipeline-design.adoc[Pipeline Design Decisions] * xref:explanation/dsl-vs-handcoded.adoc[DSL Networks vs Hand-Coded Runtimes] +* xref:explanation/bert-dsl.adoc[BERT Completely Defined in the DSL] * xref:explanation/tokenizer-internals.adoc[Tokenizer Internals] * xref:explanation/weight-quantization.adoc[Weight Quantization and Numeric Representation] * xref:explanation/embeddings.adoc[How Embeddings Work] diff --git a/docs/modules/ROOT/pages/explanation/bert-dsl.adoc b/docs/modules/ROOT/pages/explanation/bert-dsl.adoc new file mode 100644 index 00000000..affedc78 --- /dev/null +++ b/docs/modules/ROOT/pages/explanation/bert-dsl.adoc @@ -0,0 +1,144 @@ += BERT Completely Defined in the DSL +:description: How the BERT encoder is declared once in the SKaiNET NN DSL and executed eagerly, as an optimized ComputeGraph, or exported to StableHLO. + +Since 0.36.0, BERT has no hand-coded forward pass left in this repository. The entire +architecture is a single declarative definition — `bertNetwork()` — and everything else +(eager execution, graph tracing and fusion, StableHLO export, weight loading) is derived +from that one definition. This page explains what that means concretely and why the +design looks the way it does. For the general DSL-vs-hand-coded argument, see +xref:explanation/dsl-vs-handcoded.adoc[DSL Networks vs Hand-Coded Runtimes]. + +== One definition, three ways to run it + +The SKaiNET engine's core path is: *define the model once in the Kotlin DSL, then either +execute the module tree eagerly or capture it as a graph — without rewriting it.* For BERT +this is now fully realized: + +[cols="1,2"] +|=== +|Path |What happens + +|`DIRECT` (default) +|`BertEncoderRuntime` walks the module tree eagerly on the execution context — the primary +JVM path, easy to debug, no tracing involved. + +|`OPTIMIZED` +|The same module tree is traced into a ComputeGraph and run through the LLM optimization +pipeline (transpose elimination, op fusion, dead-code elimination). Graphs are +shape-specialized per sequence length and kept in a small LRU cache. Output is bit-exact +against `DIRECT`. + +|StableHLO export +|`exportTape(seqLen)` captures the encoder trace for the MLIR / StableHLO lowering shared +with the engine — `gather`, `dot_general`, and SDPA survive the lowering. Export is +gate-tested; executing the exported module (IREE) is not in scope yet. +|=== + +Because pooling and projection live *outside* the DSL network (see below), all three paths +share literally the same encoder definition. + +== The definition + +`bertNetwork()` in `llm-inference/bert` declares the whole encoder as a `sequential`: + +[source,kotlin] +---- +sequential { + // Complete embeddings block: word + position + token_type, LayerNorm + modules += BertEmbeddings(config, T::class) + + for (layer in 0 until nLayers) { + // attn block: MHA(bidirectional, bias) → Residual → LayerNorm + // ffn block: Dense → GeLU → Dense → Residual → LayerNorm + } +} +---- + +Feed it an `[L]`-shaped token-id tensor and it produces `[L, hiddenSize]` hidden states. +Two details make this definition complete and traceable where earlier attempts were not. + +=== Index-free position and token-type embeddings + +BERT sums three embeddings per token: word, absolute position, and token type. Word +embeddings use the input tensor's ids with `ops.gather` — but position ids (`0..L-1`) and +token-type ids (all zeros for sentence embedding) would normally require *synthesizing index +tensors*, which pollutes a traced graph with extra non-parameter leaves. + +`BertEmbeddings` avoids the indices entirely: + +* *Positions*: rows `0..L-1` of the position table _are_ the position vectors, in order — + so `ops.narrow(table, 0, 0, L)` selects them with no index tensor at all. +* *Token type*: sentence-embedding callers always pass segment 0, so row 0 of the type + table is reshaped to `[dim]` and broadcast-added — the same `[L, dim] + [dim]` broadcast + a Linear layer uses for bias. + +The result: a traced forward has *exactly one non-parameter leaf* — the token-id tensor. +That makes compiled-mode input detection unambiguous and the exported graph a clean +`tokens → hidden-states` encoder. The trade-off is deliberate: two-segment (cross-encoder) +inputs are not expressible, which sentence-embedding workloads never need. + +=== Post-norm residual wiring: two blocks per layer + +Decoder stacks (Llama & friends) are *pre-norm*: normalize, transform, then add the +residual. BERT is *post-norm*: + + h1 = LayerNorm(x + MHA(x)) + h2 = LayerNorm(h1 + FFN(h1)) + +The DSL's transformer blocks wire each `ResidualAdd` back to the value at the start of its +residual segment — the input of the module right after the previous `ResidualAdd`. That +rule is correct for pre-norm stacks, but in a single-block BERT layer it makes the FFN +residual grab the value *before* the post-attention LayerNorm instead of after it — a +subtle numerical bug that survives smoke tests and only shows up in parity comparisons. + +The fix is structural, not a special case: each encoder layer is defined as *two* blocks +(`encoder.layer.N.attn`, `encoder.layer.N.ffn`). Within a block, the first residual +segment starts at the block input — which places both residual boundaries exactly where +post-norm needs them. + +== What stays outside the graph — and why + +`BertEncoderRuntime` adds what sentence embedding needs *on top of* the pure encoder: + +1. *Masked mean pooling* over token positions, +2. the optional sentence-transformers `2_Dense` projection (applied even when the head is + bias-free — LEAF models ship `bias=false`, which the legacy runtime silently dropped), +3. *L2 normalization*. + +These deliberately live outside the DSL network. The pooling mask is dynamic per call, so +baking it into the graph would force retracing; and keeping the traced/exported artifact a +pure `tokens → hidden-states` encoder means the same export serves any pooling strategy a +downstream consumer picks. + +== Weight loading is derived, not written + +DSL modules have stable parameter paths (`encoder.layer.3.attn/attention/q_proj`, …). +`createBertEncoderRuntime` maps checkpoint tensors onto them via `WeightMapper` + +`BertSafeTensorsNameResolver` — there is no hand-written "load layer 3's query weight" +code left. One level up, `BertEmbeddingModel.fromHuggingFace(...)` / +`fromSafeTensors(...)` (in `llm-providers`) adds tokenizer detection, config parsing, and +Hub download, presenting the whole stack behind the neutral `EmbeddingModel` SPI. + +== What was removed, and how it was verified + +0.36.0 removes the deprecated hand-coded eager stack: `BertRuntime`, +`BertRuntimeWeights` / `BertLayerWeights`, `loadBertWeights`, `BertWeightMapper`, +`BertTensorNames`, `BertIngestion`. Migration targets: + +* `createBertEncoderRuntime(config, tensors, ctx)` — drop-in runtime level, or +* `BertEmbeddingModel.fromSafeTensors(...)` / `fromHuggingFace(...)` — the one-call level. + +The removal was gated on parity, not on review alone: + +* Hidden states within `2.2e-6` and final embeddings within `9e-8` of the + PyTorch-validated legacy runtime on real `MongoDB/mdbr-leaf-mt`; +* `DIRECT` vs `OPTIMIZED` bit-exact on synthetic and real models; +* the Java surface's reference smoke test passed *unmodified* across the swap; +* downstream, the SK-leaf CLI re-indexed its 56-chunk reference corpus with identical + embeddings — in 44.5 s instead of 676.9 s (~15×). + +== Where to go next + +* xref:tutorials/getting-started-leaf.adoc[Getting Started with LEAF] — use the result in ten lines. +* xref:tutorials/embeddings.adoc[Embeddings — Getting Started] — the `EmbeddingModel` SPI tutorial. +* xref:explanation/dsl-vs-handcoded.adoc[DSL Networks vs Hand-Coded Runtimes] — the project-wide policy this completes. diff --git a/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc b/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc index 3439cac4..0985f075 100644 --- a/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc +++ b/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc @@ -110,6 +110,10 @@ The goal is to extend the DSL to support these patterns over time. |`bertNetwork()` |DSL-only: `BertEncoderRuntime` (eager + traced graph); legacy runtime removed +|T5 / GTR (vec2text) +|_none_ +|Hand-coded encoder-decoder (`T5Runtime`); relative-position bias and the per-step decode loop are not DSL-expressible yet + |Voxtral |`voxtralBackboneNetwork()` |Partial DSL diff --git a/docs/modules/ROOT/pages/index.adoc b/docs/modules/ROOT/pages/index.adoc index b8a63d02..1cb1b14f 100644 --- a/docs/modules/ROOT/pages/index.adoc +++ b/docs/modules/ROOT/pages/index.adoc @@ -43,6 +43,11 @@ It loads GGUF and SafeTensors models, builds compute graphs from DSL network def |No |`bertNetwork()` +|T5 +|gtr-t5-base (GTR embedder, vec2text inversion) +|No +|Hand-coded + |Voxtral |Voxtral TTS |No diff --git a/docs/modules/ROOT/pages/tutorials/embeddings.adoc b/docs/modules/ROOT/pages/tutorials/embeddings.adoc index f0282888..8c707910 100644 --- a/docs/modules/ROOT/pages/tutorials/embeddings.adoc +++ b/docs/modules/ROOT/pages/tutorials/embeddings.adoc @@ -113,5 +113,7 @@ For the BERT entry, the script computes embeddings for the prompt and the docume == What's Next +* xref:tutorials/getting-started-leaf.adoc[Getting Started with LEAF] — a model-focused walkthrough with a complete semantic-search example. +* xref:explanation/bert-dsl.adoc[BERT Completely Defined in the DSL] — how the encoder behind this API is defined, executed, and exported. * xref:explanation/embeddings.adoc[How embeddings work — pooling, normalization, dimensionality] — the deeper "why". * xref:tutorials/getting-started-java.adoc[Getting Started for Java Developers] — the analogous Java surface for chat / tool calling. diff --git a/docs/modules/ROOT/pages/tutorials/getting-started-leaf.adoc b/docs/modules/ROOT/pages/tutorials/getting-started-leaf.adoc new file mode 100644 index 00000000..0326877c --- /dev/null +++ b/docs/modules/ROOT/pages/tutorials/getting-started-leaf.adoc @@ -0,0 +1,156 @@ += Getting Started with LEAF +:description: Load MongoDB's LEAF embedding models in one call and build semantic search on the JVM — no Python, no server. + +https://huggingface.co/MongoDB/mdbr-leaf-ir[LEAF^] is MongoDB's family of compact BERT-style +embedding models for retrieval workloads. At ~23M parameters (6 layers, hidden size 384) a LEAF +model loads in seconds, runs comfortably on CPU, and is small enough to embed in a desktop app, +a CLI, or an Android service — which makes it a natural first model for SKaiNET Transformers' +embedding stack. + +Two variants are published: + +[cols="1,1,2"] +|=== +|Model |Output dimensions |Tuned for + +|`MongoDB/mdbr-leaf-ir` +|768 +|Information retrieval — query/document search, RAG indexing + +|`MongoDB/mdbr-leaf-mt` +|1024 +|Multi-task — retrieval plus classification, clustering, similarity +|=== + +Both ship a bias-free sentence-transformers `2_Dense` projection head that maps the 384-dim +encoder output to the advertised dimensionality. The runtime applies it automatically (since +0.36.0 — earlier releases silently skipped bias-free heads and returned raw 384-dim vectors). + +== Prerequisites + +* JDK 21+ (Java 25 preferred for the Vector API) +* The `llm-providers` artifact (via the BOM): + +[source,kotlin] +---- +dependencies { + implementation(platform("sk.ainet.transformers:skainet-transformers-bom:0.36.0")) + implementation("sk.ainet.transformers:skainet-transformers-providers") +} +---- + +No model setup is needed — the factory downloads the snapshot from the Hugging Face Hub on +first use and caches it under `~/.cache/skainet/models/` (offline-safe afterwards; set +`HF_TOKEN` for gated repos). + +== One call to a model + +[source,kotlin] +---- +import sk.ainet.llm.providers.BertEmbeddingModel + +BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-ir").use { model -> + println(model.dimensions) // 768 + val v: FloatArray = model.embed("What is a pangram?") +} +---- + +The returned `EmbeddingModel` (the provider-neutral SPI from `llm-api`) hides the whole stack: +tokenizer, the DSL-defined encoder (`bertNetwork()`), masked mean pooling, the `2_Dense` +projection, and L2 normalization. Vectors come out unit-length, so cosine similarity is just a +dot product. + +Already have a local snapshot (for example from `huggingface-cli download`)? Point at the +directory instead — `config.json`, the tokenizer (`vocab.txt` or `tokenizer.json`), +`model.safetensors`, and the optional `2_Dense/` head are auto-detected: + +[source,kotlin] +---- +val model = BertEmbeddingModel.fromSafeTensors(Path.of("/models/mdbr-leaf-ir")) +---- + +== Semantic search in 30 lines + +Index a handful of documents, then rank them against a query: + +[source,kotlin] +---- +import sk.ainet.llm.providers.BertEmbeddingModel + +fun dot(a: FloatArray, b: FloatArray): Float { + var s = 0f + for (i in a.indices) s += a[i] * b[i] + return s // vectors are L2-normalized — this IS the cosine similarity +} + +fun main() { + BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-ir").use { model -> + val documents = listOf( + "A pangram is a sentence that contains every letter of the alphabet.", + "The Eiffel Tower is located on the Champ de Mars in Paris.", + "Kotlin Multiplatform shares code between JVM, Native, and JS targets.", + ) + + // Index: one batched call, response order matches request order. + val index: List = model.embed(documents) + + // Query + val query = model.embed("sentence with all 26 letters") + documents.zip(index) + .sortedByDescending { (_, v) -> dot(query, v) } + .forEach { (doc, v) -> println("%.4f %s".format(dot(query, v), doc)) } + } +} +---- + +The pangram document wins by a wide margin. Swap the in-memory list for a vector store and this +is the indexing side of a RAG pipeline — see +https://github.com/michalharakal/SK-leaf[SK-leaf^] for a complete CLI (index + ask) built on +exactly this API. + +== From Java + +`BertEmbeddingModel` is `@JvmStatic` throughout: + +[source,java] +---- +import sk.ainet.llm.api.EmbeddingModel; +import sk.ainet.llm.providers.BertEmbeddingModel; + +try (EmbeddingModel model = BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-ir")) { + float[] vector = model.embed("The quick brown fox"); +} +---- + +For a session-style surface (`encode` / `similarity`) see `KBertJava` in +xref:tutorials/embeddings.adoc[Embeddings — Getting Started]. + +== From the CLI + +`kbert-cli` accepts a Hub repo id (or a local snapshot directory) directly: + +[source,bash] +---- +./gradlew :llm-apps:kbert-cli:run \ + --args="MongoDB/mdbr-leaf-ir 'pangram' 'A pangram is a sentence that contains every letter of the alphabet.'" +---- + +It prints the embeddings' dimensionality and the query/document cosine similarity. + +== Performance notes + +* The default execution mode is `DIRECT` (eager) — the primary, well-trodden JVM path. As of + 0.36.0 it runs the DSL-defined encoder; indexing the SK-leaf reference corpus (56 chunks) + dropped from 676.9 s to 44.5 s (~15×) versus the removed legacy runtime, with identical + embeddings. +* `BertEncoderRuntime` also offers an `OPTIMIZED` mode that executes a traced, fused + ComputeGraph (bit-exact vs `DIRECT`, shape-specialized per sequence length) — see + xref:explanation/bert-dsl.adoc[BERT Completely Defined in the DSL]. +* For the fastest matmuls on the JVM, add the native FFM CPU backend + (`sk.ainet.core:skainet-backend-native-cpu`) — it is auto-discovered when present. + +== What's Next + +* xref:tutorials/embeddings.adoc[Embeddings — Getting Started] — the SPI-level tutorial (batching, `EmbeddingRequest`, Java session facade). +* xref:explanation/bert-dsl.adoc[BERT Completely Defined in the DSL] — how the encoder is built, executed, and exported. +* xref:explanation/embeddings.adoc[How Embeddings Work] — pooling, normalization, dimensionality. diff --git a/gradle.properties b/gradle.properties index 99e93ccb..5ecbcf49 100644 --- a/gradle.properties +++ b/gradle.properties @@ -1,5 +1,5 @@ GROUP=sk.ainet.transformers -VERSION_NAME=0.35.0 +VERSION_NAME=0.36.0 POM_DESCRIPTION=SKaiNET-transformers diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index f94a3aff..f0980ba1 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -1,5 +1,5 @@ [versions] -skainet = "0.35.0" +skainet = "0.36.0" agp = "9.2.1" jacksonDatabind = "2.22.1" jsonSchemaValidator = "3.0.6" diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/graph/LLMFusedOpHandlers.kt b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/graph/LLMFusedOpHandlers.kt index 5f8c5922..e5400407 100644 --- a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/graph/LLMFusedOpHandlers.kt +++ b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/graph/LLMFusedOpHandlers.kt @@ -20,39 +20,6 @@ public object LLMFusedOpHandlers { ComputeGraphExecutor.registerFusedOp("fused_rms_norm", RmsNormHandler) ComputeGraphExecutor.registerFusedOp("fused_swiglu_ffn", SwiGluFFNHandler) ComputeGraphExecutor.registerFusedOp("fused_qkv_proj", QKVProjHandler) - ComputeGraphExecutor.registerFusedOp("permute", PermuteHandler) - } - - /** - * Replays `permute` with its recorded `axes` parameter. - * - * Works around the executor's builtin dispatch (engine ≤ 0.35.0), which - * replays "permute" as a plain last-two-dims [TensorOps.transpose] and - * silently drops the axes — wrong for any rank-3 permutation like MHA's - * heads/sequence swap `[1, 0, 2]` on multi-token (encoder / prefill) - * traces. Single-token decode never hits it (the seqLen-1 shortcut - * reshapes instead), which is why generation paths didn't surface this. - * Remove once the engine's builtin honors `axes`. - */ - private object PermuteHandler : FusedOpHandler { - override fun execute( - ops: TensorOps, - inputs: List>, - params: Map - ): List> { - val axesParam = params["axes"] ?: params["dims"] - val axes = when (axesParam) { - is IntArray -> axesParam - is List<*> -> axesParam.mapNotNull { (it as? Number)?.toInt() }.toIntArray() - else -> null - } - return if (axes != null && axes.size == inputs[0].rank) { - listOf(ops.permute(inputs[0], axes)) - } else { - // No usable axes recorded — preserve the executor's legacy behavior. - listOf(ops.transpose(inputs[0])) - } - } } private object RmsNormHandler : FusedOpHandler { 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 a3906f96..1c4df175 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 @@ -187,8 +187,7 @@ public class BertEncoderRuntime( } private fun compiledFor(sampleTokenIds: IntArray, diagnostics: MutableList): CompiledEncoder { - // Correct replay handlers (incl. the axes-aware permute override the - // multi-token attention trace depends on). + // Fused-op replay handlers (RMSNorm / SwiGLU / QKV decompositions). LLMFusedOpHandlers.registerAll() // 1. Record one eager forward through a tape-recording context. diff --git a/llm-inference/t5/src/commonMain/kotlin/sk/ainet/models/t5/T5Runtime.kt b/llm-inference/t5/src/commonMain/kotlin/sk/ainet/models/t5/T5Runtime.kt index ac46e398..808a4843 100644 --- a/llm-inference/t5/src/commonMain/kotlin/sk/ainet/models/t5/T5Runtime.kt +++ b/llm-inference/t5/src/commonMain/kotlin/sk/ainet/models/t5/T5Runtime.kt @@ -21,8 +21,8 @@ import sk.ainet.lang.types.DType import kotlin.reflect.KClass /** - * Hand-coded T5 encoder-decoder runtime, written in the direct tensor-ops style of - * [sk.ainet.models.bert.BertRuntime] (per-head attention via narrow/matmul/softmax; no + * Hand-coded T5 encoder-decoder runtime, written in the direct tensor-ops style of the + * (since-removed) hand-coded BERT runtime (per-head attention via narrow/matmul/softmax; no * Module/graph composition, no KV cache). Batch size 1. * * T5 specifics vs. a vanilla transformer, all handled here: