diff --git a/CHANGELOG.md b/CHANGELOG.md index d556b688..fdcb5c3f 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index f349c5cd..5d4c6291 100644 --- a/README.md +++ b/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 diff --git a/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc b/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc index 64f39fb0..3439cac4 100644 --- a/docs/modules/ROOT/pages/explanation/dsl-vs-handcoded.adoc +++ b/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()` diff --git a/docs/modules/ROOT/pages/explanation/embeddings.adoc b/docs/modules/ROOT/pages/explanation/embeddings.adoc index 8c4626e1..6f3ba317 100644 --- a/docs/modules/ROOT/pages/explanation/embeddings.adoc +++ b/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,7 +72,7 @@ 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 @@ -80,8 +80,14 @@ The same text can produce different vectors with different tokenizers — this i |=== | 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` diff --git a/docs/modules/ROOT/pages/tutorials/embeddings.adoc b/docs/modules/ROOT/pages/tutorials/embeddings.adoc index 610ee725..f0282888 100644 --- a/docs/modules/ROOT/pages/tutorials/embeddings.adoc +++ b/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(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"); } ---- diff --git a/docs/specs/spring-ai-adapter.md b/docs/specs/spring-ai-adapter.md index 36329df5..dc3cf7c7 100644 --- a/docs/specs/spring-ai-adapter.md +++ b/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` (wraps any `InferenceRuntime` + `Tokenizer` + `ChatTemplate`) and `SkaiNetEmbeddingModel` - (wraps `BertRuntime`). + (wraps `BertEncoderRuntime`). - 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 diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index d4d49f8f..5a87535e 100644 --- a/gradle/libs.versions.toml +++ b/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" } diff --git a/llm-apps/kbert-cli/build.gradle.kts b/llm-apps/kbert-cli/build.gradle.kts index e083ce09..952c068d 100644 --- a/llm-apps/kbert-cli/build.gradle.kts +++ b/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) diff --git a/llm-apps/kbert-cli/src/main/kotlin/sk/ainet/apps/bert/cli/Demo.kt b/llm-apps/kbert-cli/src/main/kotlin/sk/ainet/apps/bert/cli/Demo.kt index 7fca96af..170901b5 100644 --- a/llm-apps/kbert-cli/src/main/kotlin/sk/ainet/apps/bert/cli/Demo.kt +++ b/llm-apps/kbert-cli/src/main/kotlin/sk/ainet/apps/bert/cli/Demo.kt @@ -1,174 +1,126 @@ package sk.ainet.apps.bert.cli -import sk.ainet.models.bert.BertRuntime -import sk.ainet.models.bert.HuggingFaceTokenizer -import sk.ainet.models.bert.loadBertWeights -import sk.ainet.context.DirectCpuExecutionContext -import sk.ainet.io.JvmRandomAccessSource -import sk.ainet.io.safetensors.SafeTensorsParametersLoader -import sk.ainet.lang.types.FP32 -import kotlinx.coroutines.runBlocking -import java.nio.file.Path +import sk.ainet.llm.api.EmbeddingModel import java.util.Locale -import kotlin.io.path.exists -import kotlin.io.path.readText import kotlin.system.exitProcess import kotlin.time.measureTime private fun usage(): Nothing { - println("Usage: kbert-demo [--iterations N]") + println("Usage: kbert-demo [--iterations N]") println() println(" Directory containing model.safetensors, vocab.txt, config.json") + println(" Hugging Face repo — downloaded and cached on first use") println(" --iterations N Number of embeddings for throughput test (default: 100)") exitProcess(1) } fun main(args: Array) { - runBlocking { - if (args.isEmpty()) usage() - - val modelDir = Path.of(args[0]) - if (!modelDir.exists()) error("Model directory not found: $modelDir") - - var iterations = 100 - var i = 1 - while (i < args.size) { - when (args[i]) { - "--iterations" -> { - iterations = args.getOrNull(i + 1)?.toIntOrNull() - ?: error("--iterations requires an integer argument") - i += 2 - } - else -> error("Unknown flag: ${args[i]}") - } - } - - // Detect config - val config = detectConfig(modelDir) - - println("=== SKaiNET LEAF IR Demo ===") - println("Model: ${modelDir.fileName}") - println("Config: hidden=${config.hiddenSize}, layers=${config.numHiddenLayers}, heads=${config.numAttentionHeads}, projection=${config.projectionDim ?: "none"}") - println() - - // Load tokenizer - val vocabPath = modelDir.resolve("vocab.txt") - if (!vocabPath.exists()) error("vocab.txt not found in $modelDir") - print("Loading tokenizer... ") - val tokenizer = HuggingFaceTokenizer.fromVocabTxt(vocabPath.readText()) - println("done (vocab=${tokenizer.vocabSize})") - - // Load model weights using multi-loader to pick up both model.safetensors and 2_Dense/model.safetensors - print("Loading model weights... ") - val ctx = DirectCpuExecutionContext() - val loaders = buildList { - val mainFile = resolveModelFile(modelDir) - add(SafeTensorsParametersLoader( - sourceProvider = { JvmRandomAccessSource.open(mainFile.toString()) }, - onProgress = { _, _, _ -> } - )) - val denseFile = modelDir.resolve("2_Dense/model.safetensors") - if (denseFile.exists()) { - add(SafeTensorsParametersLoader( - sourceProvider = { JvmRandomAccessSource.open(denseFile.toString()) }, - onProgress = { _, _, _ -> } - )) + if (args.isEmpty()) usage() + + val modelRef = args[0] + var iterations = 100 + var i = 1 + while (i < args.size) { + when (args[i]) { + "--iterations" -> { + iterations = args.getOrNull(i + 1)?.toIntOrNull() + ?: error("--iterations requires an integer argument") + i += 2 } + else -> error("Unknown flag: ${args[i]}") } - val weights = loadBertWeights(loaders, ctx, FP32::class, config) - println("done (${loaders.size} file(s))") - - val runtime = BertRuntime(ctx, weights, FP32::class) - val expectedDim = config.projectionDim ?: config.hiddenSize + } - // Helper: encode text and return float list - fun encode(text: String): List { - val tok = tokenizer.encodeWithMetadata(text) - return runtime.encode(tok.inputIds, tok.attentionMask, tok.tokenTypeIds).expectFloatBuffer() - } + println("=== SKaiNET LEAF IR Demo ===") + print("Loading model $modelRef... ") + val model: EmbeddingModel + val loadElapsed = measureTime { model = loadModel(modelRef) } + println("done ($loadElapsed)") + println("Embedding dimensions: ${model.dimensions}") - // ── Section 1: Single Embedding Sanity Check ── - println() - println("=== 1. Single Embedding Sanity Check ===") - val query1 = "What is artificial intelligence?" - val emb1 = encode(query1) - println("Query: \"$query1\"") - println("Embedding dim: ${emb1.size} (expected: $expectedDim)") - if (emb1.size != expectedDim) { - println("FAIL: dimension mismatch!") - exitProcess(1) - } - val nonZero = emb1.count { it != 0f } - val allFinite = emb1.all { it.isFinite() } - val min = emb1.min() - val max = emb1.max() - val mean = emb1.map { it.toDouble() }.average() - println(String.format(Locale.ROOT, "Non-zero: %d/%d, All finite: %s", nonZero, emb1.size, allFinite)) - println(String.format(Locale.ROOT, "Range: [%.6f, %.6f], Mean: %.6f", min, max, mean)) - println("PASS") + fun encode(text: String): FloatArray = model.embed(text) - // ── Section 2: Semantic Similarity Check ── - println() - println("=== 2. Semantic Similarity Check ===") - val query2 = "Define artificial intelligence" - val query3 = "What is the weather today?" - val emb2 = encode(query2) - val emb3 = encode(query3) - val sim12 = cosineSimilarity(emb1, emb2) - val sim13 = cosineSimilarity(emb1, emb3) - println(String.format(Locale.ROOT, "sim(\"%s\", \"%s\") = %.4f", query1, query2, sim12)) - println(String.format(Locale.ROOT, "sim(\"%s\", \"%s\") = %.4f", query1, query3, sim13)) - if (sim12 > sim13) { - println("PASS: related pair scores higher than unrelated pair") - } else { - println("FAIL: expected related pair to score higher") - } + // ── Section 1: Single Embedding Sanity Check ── + println() + println("=== 1. Single Embedding Sanity Check ===") + val query1 = "What is artificial intelligence?" + val emb1 = encode(query1) + println("Query: \"$query1\"") + println("Embedding dim: ${emb1.size} (expected: ${model.dimensions})") + if (emb1.size != model.dimensions) { + println("FAIL: dimension mismatch!") + exitProcess(1) + } + val nonZero = emb1.count { it != 0f } + val allFinite = emb1.all { it.isFinite() } + val min = emb1.min() + val max = emb1.max() + val mean = emb1.map { it.toDouble() }.average() + println(String.format(Locale.ROOT, "Non-zero: %d/%d, All finite: %s", nonZero, emb1.size, allFinite)) + println(String.format(Locale.ROOT, "Range: [%.6f, %.6f], Mean: %.6f", min, max, mean)) + println("PASS") + + // ── Section 2: Semantic Similarity Check ── + println() + println("=== 2. Semantic Similarity Check ===") + val query2 = "Define artificial intelligence" + val query3 = "What is the weather today?" + val emb2 = encode(query2) + val emb3 = encode(query3) + val sim12 = cosineSimilarity(emb1, emb2) + val sim13 = cosineSimilarity(emb1, emb3) + println(String.format(Locale.ROOT, "sim(\"%s\", \"%s\") = %.4f", query1, query2, sim12)) + println(String.format(Locale.ROOT, "sim(\"%s\", \"%s\") = %.4f", query1, query3, sim13)) + if (sim12 > sim13) { + println("PASS: related pair scores higher than unrelated pair") + } else { + println("FAIL: expected related pair to score higher") + } - // ── Section 3: Small IR Retrieval Test ── - println() - println("=== 3. Small IR Retrieval Test ===") - val baseQuery = "MongoDB is a NoSQL database" - val candidates = listOf( - "MongoDB stores data in documents", - "PostgreSQL is a relational database", - "The cat sat on the mat", - "NoSQL databases are non-relational", - "MongoDB uses BSON format" - ) - val baseEmb = encode(baseQuery) - println("Base query: \"$baseQuery\"") - println() + // ── Section 3: Small IR Retrieval Test ── + println() + println("=== 3. Small IR Retrieval Test ===") + val baseQuery = "MongoDB is a NoSQL database" + val candidates = listOf( + "MongoDB stores data in documents", + "PostgreSQL is a relational database", + "The cat sat on the mat", + "NoSQL databases are non-relational", + "MongoDB uses BSON format" + ) + val baseEmb = encode(baseQuery) + println("Base query: \"$baseQuery\"") + println() - val scores = candidates.mapIndexed { idx, doc -> - val docEmb = encode(doc) - val sim = cosineSimilarity(baseEmb, docEmb) - println(String.format(Locale.ROOT, " [%d] sim=%.4f %s", idx, sim, doc)) - idx to sim - } - val ranked = scores.sortedByDescending { it.second } - val bestIdx = ranked.first().first - println() - println(String.format(Locale.ROOT, "Best match: [%d] score=%.4f", bestIdx, ranked.first().second)) - val mongoRelated = setOf(0, 3, 4) - if (bestIdx in mongoRelated) { - println("PASS: best match is MongoDB-related") - } else { - println("FAIL: expected a MongoDB-related doc to rank first") - } + val scores = candidates.mapIndexed { idx, doc -> + val docEmb = encode(doc) + val sim = cosineSimilarity(baseEmb, docEmb) + println(String.format(Locale.ROOT, " [%d] sim=%.4f %s", idx, sim, doc)) + idx to sim + } + val ranked = scores.sortedByDescending { it.second } + val bestIdx = ranked.first().first + println() + println(String.format(Locale.ROOT, "Best match: [%d] score=%.4f", bestIdx, ranked.first().second)) + val mongoRelated = setOf(0, 3, 4) + if (bestIdx in mongoRelated) { + println("PASS: best match is MongoDB-related") + } else { + println("FAIL: expected a MongoDB-related doc to rank first") + } - // ── Section 4: Performance / Throughput Test ── - println() - println("=== 4. Performance / Throughput Test ($iterations iterations) ===") - val elapsed = measureTime { - for (j in 0 until iterations) { - encode("$query1 $j") - } + // ── Section 4: Performance / Throughput Test ── + println() + println("=== 4. Performance / Throughput Test ($iterations iterations) ===") + val elapsed = measureTime { + for (j in 0 until iterations) { + encode("$query1 $j") } - val avgMs = elapsed.inWholeNanoseconds / 1_000_000.0 / iterations - println(String.format(Locale.ROOT, "Total: %s for %d embeddings", elapsed, iterations)) - println(String.format(Locale.ROOT, "Avg: %.3f ms/embedding", avgMs)) - - println() - println("=== Demo Complete ===") } + val avgMs = elapsed.inWholeNanoseconds / 1_000_000.0 / iterations + println(String.format(Locale.ROOT, "Total: %s for %d embeddings", elapsed, iterations)) + println(String.format(Locale.ROOT, "Avg: %.3f ms/embedding", avgMs)) + + println() + println("=== Demo Complete ===") } diff --git a/llm-apps/kbert-cli/src/main/kotlin/sk/ainet/apps/bert/cli/Main.kt b/llm-apps/kbert-cli/src/main/kotlin/sk/ainet/apps/bert/cli/Main.kt index b80a3063..83d30215 100644 --- a/llm-apps/kbert-cli/src/main/kotlin/sk/ainet/apps/bert/cli/Main.kt +++ b/llm-apps/kbert-cli/src/main/kotlin/sk/ainet/apps/bert/cli/Main.kt @@ -1,158 +1,61 @@ package sk.ainet.apps.bert.cli -import sk.ainet.models.bert.BertIngestion -import sk.ainet.models.bert.BertModelConfig -import sk.ainet.models.bert.BertRuntime -import sk.ainet.models.bert.HuggingFaceTokenizer -import sk.ainet.models.bert.MDBR_LEAF_IR_CONFIG -import sk.ainet.models.bert.BertRuntimeWeights -import sk.ainet.context.DirectCpuExecutionContext -import sk.ainet.io.JvmRandomAccessSource -import sk.ainet.io.safetensors.SafeTensorsParametersLoader -import sk.ainet.lang.tensor.data.FloatArrayTensorData -import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.types.DType -import sk.ainet.lang.types.FP32 -import kotlinx.coroutines.runBlocking +import sk.ainet.llm.api.EmbeddingModel +import sk.ainet.llm.providers.BertEmbeddingModel import java.nio.file.Path -import kotlin.io.path.exists -import kotlin.io.path.readText +import kotlin.io.path.isDirectory import kotlin.system.exitProcess import kotlin.time.measureTime private fun usage(): Nothing { - println("Usage: kbert \"query text\" [\"doc text\"]") + println("Usage: kbert \"query text\" [\"doc text\"]") println() println(" Directory containing model.safetensors, vocab.txt, config.json") + println(" Hugging Face repo (e.g. MongoDB/mdbr-leaf-mt) — downloaded and") + println(" cached under ~/.cache/skainet/models on first use") println(" \"query text\" Text to encode") println(" \"doc text\" Optional second text — if given, prints cosine similarity") exitProcess(1) } fun main(args: Array) { - runBlocking { - if (args.isEmpty()) usage() - - val modelDir = Path.of(args[0]) - if (!modelDir.exists()) error("Model directory not found: $modelDir") - - val textA = args.getOrNull(1) ?: usage() - val textB = args.getOrNull(2) - - // Resolve model files - val safetensorsPath = resolveModelFile(modelDir) - val vocabPath = modelDir.resolve("vocab.txt") - if (!vocabPath.exists()) error("vocab.txt not found in $modelDir") - - // Detect config - val config = detectConfig(modelDir) - - println("Model: ${modelDir.fileName}") - println("Config: hidden=${config.hiddenSize}, layers=${config.numHiddenLayers}, heads=${config.numAttentionHeads}") - println() - - // Load tokenizer - print("Loading tokenizer... ") - val tokenizer = HuggingFaceTokenizer.fromVocabTxt(vocabPath.readText()) - println("done (vocab=${tokenizer.vocabSize})") - - // Load model weights - print("Loading model weights... ") - val ctx = DirectCpuExecutionContext() - val ingestion = BertIngestion(ctx, FP32::class, config) - val loader = SafeTensorsParametersLoader( - sourceProvider = { JvmRandomAccessSource.open(safetensorsPath.toString()) }, - onProgress = { _, _, _ -> } - ) - val weights: BertRuntimeWeights - val loadElapsed = measureTime { weights = ingestion.load(loader) } - println("done (${loadElapsed})") - - val runtime = BertRuntime(ctx, weights, FP32::class) - - // Encode text(s) - val tokOutputA = tokenizer.encodeWithMetadata(textA) - println("\nEncoding: \"$textA\" (${tokOutputA.inputIds.size} tokens)") - val embA: Tensor - val encodeElapsed = measureTime { - embA = runtime.encode(tokOutputA.inputIds, tokOutputA.attentionMask, tokOutputA.tokenTypeIds) - } - println("Encoded in $encodeElapsed") - val vecA = embA.expectFloatBuffer() - println("Embedding (first 8): ${vecA.take(8).joinToString(", ") { "%.6f".format(it) }}") - println("Embedding dim: ${vecA.size}") - - if (textB != null) { - val tokOutputB = tokenizer.encodeWithMetadata(textB) - println("\nEncoding: \"$textB\" (${tokOutputB.inputIds.size} tokens)") - val embB = runtime.encode(tokOutputB.inputIds, tokOutputB.attentionMask, tokOutputB.tokenTypeIds) - val vecB = embB.expectFloatBuffer() - println("Embedding (first 8): ${vecB.take(8).joinToString(", ") { "%.6f".format(it) }}") - - // Cosine similarity - val sim = cosineSimilarity(vecA, vecB) - println("\nCosine similarity: %.6f".format(sim)) - } + if (args.isEmpty()) usage() + + val modelRef = args[0] + val textA = args.getOrNull(1) ?: usage() + val textB = args.getOrNull(2) + + print("Loading model $modelRef... ") + val model: EmbeddingModel + val loadElapsed = measureTime { model = loadModel(modelRef) } + println("done ($loadElapsed) — dimensions=${model.dimensions}") + + println("\nEncoding: \"$textA\"") + val vecA: FloatArray + val encodeElapsed = measureTime { vecA = model.embed(textA) } + println("Encoded in $encodeElapsed") + println("Embedding (first 8): ${vecA.take(8).joinToString(", ") { "%.6f".format(it) }}") + println("Embedding dim: ${vecA.size}") + + if (textB != null) { + println("\nEncoding: \"$textB\"") + val vecB = model.embed(textB) + println("Embedding (first 8): ${vecB.take(8).joinToString(", ") { "%.6f".format(it) }}") + println("\nCosine similarity: %.6f".format(cosineSimilarity(vecA, vecB))) } } -internal fun resolveModelFile(modelDir: Path): Path { - val candidates = listOf("model.safetensors", "pytorch_model.safetensors") - for (name in candidates) { - val p = modelDir.resolve(name) - if (p.exists()) return p +/** Local directory → fromSafeTensors; anything owner/name-shaped → Hugging Face. */ +internal fun loadModel(modelRef: String): EmbeddingModel { + val asPath = Path.of(modelRef) + if (asPath.isDirectory()) return BertEmbeddingModel.fromSafeTensors(asPath) + require(Regex("[\\w.-]+/[\\w.-]+").matches(modelRef)) { + "Not a directory and not an owner/name Hugging Face repo id: $modelRef" } - // Try any .safetensors file - val dir = modelDir.toFile() - val found = dir.listFiles()?.firstOrNull { it.extension == "safetensors" } - if (found != null) return found.toPath() - error("No .safetensors file found in $modelDir") + return BertEmbeddingModel.fromHuggingFace(modelRef) } -internal fun detectConfig(modelDir: Path): BertModelConfig { - val configPath = modelDir.resolve("config.json") - if (!configPath.exists()) { - println("No config.json found, using MDBR_LEAF_IR_CONFIG defaults") - return MDBR_LEAF_IR_CONFIG - } - val json = configPath.readText() - // Check 2_Dense/config.json for projection dim (sentence-transformers layout) - val denseConfigPath = modelDir.resolve("2_Dense/config.json") - val denseJson = if (denseConfigPath.exists()) denseConfigPath.readText() else null - return parseConfigJson(json, denseJson) -} - -internal fun parseConfigJson(json: String, denseJson: String? = null): BertModelConfig { - fun extractInt(source: String, key: String, default: Int): Int { - val pattern = Regex("\"$key\"\\s*:\\s*(\\d+)") - return pattern.find(source)?.groupValues?.get(1)?.toIntOrNull() ?: default - } - fun extractDouble(source: String, key: String, default: Double): Double { - val pattern = Regex("\"$key\"\\s*:\\s*([\\d.eE\\-+]+)") - return pattern.find(source)?.groupValues?.get(1)?.toDoubleOrNull() ?: default - } - - // Projection dim comes from 2_Dense/config.json (sentence-transformers layout) - val projDim = if (denseJson != null) { - extractInt(denseJson, "out_features", 0).let { if (it > 0) it else null } - } else { - null - } - - return BertModelConfig( - vocabSize = extractInt(json, "vocab_size", 30522), - hiddenSize = extractInt(json, "hidden_size", 384), - numHiddenLayers = extractInt(json, "num_hidden_layers", 6), - numAttentionHeads = extractInt(json, "num_attention_heads", 12), - intermediateSize = extractInt(json, "intermediate_size", 1536), - maxPositionEmbeddings = extractInt(json, "max_position_embeddings", 512), - typeVocabSize = extractInt(json, "type_vocab_size", 2), - layerNormEps = extractDouble(json, "layer_norm_eps", 1e-12), - projectionDim = projDim - ) -} - -internal fun cosineSimilarity(a: List, b: List): Float { +internal fun cosineSimilarity(a: FloatArray, b: FloatArray): Float { require(a.size == b.size) { "Vectors must have same dimension" } var dot = 0f var normA = 0f @@ -165,9 +68,3 @@ internal fun cosineSimilarity(a: List, b: List): Float { val denom = kotlin.math.sqrt(normA) * kotlin.math.sqrt(normB) return if (denom > 0f) dot / denom else 0f } - -internal fun Tensor.expectFloatBuffer(): List { - val data = this.data - if (data is FloatArrayTensorData<*>) return data.buffer.toList() - return data.copyToFloatArray().toList() -} 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 e5400407..5f8c5922 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,6 +20,39 @@ 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-core/src/commonMain/kotlin/sk/ainet/apps/llm/weights/LLMWeightNameResolvers.kt b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/weights/LLMWeightNameResolvers.kt index f32dc849..65159aa6 100644 --- a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/weights/LLMWeightNameResolvers.kt +++ b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/weights/LLMWeightNameResolvers.kt @@ -159,7 +159,9 @@ public class BertSafeTensorsNameResolver : WeightNameResolver { val pathParts = modulePath.split("/").drop(1) val moduleName = pathParts.lastOrNull() ?: return null val layerPart = pathParts.firstOrNull { it.startsWith("encoder.layer.") } - val layerNum = layerPart?.removePrefix("encoder.layer.")?.toIntOrNull() + // Layer blocks may carry a sub-block suffix ("encoder.layer.3.attn") — + // the layer number is the first segment after the prefix. + val layerNum = layerPart?.removePrefix("encoder.layer.")?.substringBefore('.')?.toIntOrNull() val layerPrefix = if (layerNum != null) "bert.encoder.layer.$layerNum" else null val inEmbeddings = pathParts.any { it == "embeddings" } @@ -167,6 +169,13 @@ public class BertSafeTensorsNameResolver : WeightNameResolver { moduleName == "word_embeddings" && inEmbeddings -> "bert.embeddings.word_embeddings.weight" + // BertEmbeddings' own additive tables: the module path ends at + // "embeddings"; the param name carries the table identity. + paramName.endsWith("position_embeddings.weight") && inEmbeddings -> + "bert.embeddings.position_embeddings.weight" + paramName.endsWith("token_type_embeddings.weight") && inEmbeddings -> + "bert.embeddings.token_type_embeddings.weight" + moduleName == "LayerNorm" && inEmbeddings && paramName.endsWith(".weight") -> "bert.embeddings.LayerNorm.weight" moduleName == "LayerNorm" && inEmbeddings && paramName.endsWith(".bias") -> diff --git a/llm-inference/bert/api/android/bert.api b/llm-inference/bert/api/android/bert.api index bacefc17..6306ef75 100644 --- a/llm-inference/bert/api/android/bert.api +++ b/llm-inference/bert/api/android/bert.api @@ -1,47 +1,40 @@ -public final class sk/ainet/models/bert/BertIngestion { - public fun (Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/models/bert/BertModelConfig;)V - public final fun load (Lsk/ainet/io/ParametersLoader;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; +public final class sk/ainet/models/bert/BertConfigKt { + public static final fun getMDBR_LEAF_IR_CONFIG ()Lsk/ainet/models/bert/BertModelConfig; } -public final class sk/ainet/models/bert/BertLayerWeights { - public fun (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V - public final fun component1 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component10 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component11 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component12 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component13 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component14 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component15 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component16 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component2 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component3 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component4 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component5 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component6 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component7 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component8 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component9 ()Lsk/ainet/lang/tensor/Tensor; - public final fun copy (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/models/bert/BertLayerWeights; - public static synthetic fun copy$default (Lsk/ainet/models/bert/BertLayerWeights;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;ILjava/lang/Object;)Lsk/ainet/models/bert/BertLayerWeights; - public fun equals (Ljava/lang/Object;)Z - public final fun getAttnLayerNormBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getAttnLayerNormWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getAttnOutputBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getAttnOutputWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getIntermediateBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getIntermediateWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getKeyBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getKeyWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getOutputBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getOutputLayerNormBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getOutputLayerNormWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getOutputWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getQueryBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getQueryWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getValueBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getValueWeight ()Lsk/ainet/lang/tensor/Tensor; - public fun hashCode ()I - public fun toString ()Ljava/lang/String; +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 class sk/ainet/models/bert/BertEmbeddings : sk/ainet/lang/nn/Module, sk/ainet/lang/nn/topology/ModuleParameters { + public fun (Lsk/ainet/models/bert/BertModelConfig;Lkotlin/reflect/KClass;Ljava/lang/String;)V + public synthetic fun (Lsk/ainet/models/bert/BertModelConfig;Lkotlin/reflect/KClass;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun getModules ()Ljava/util/List; + public fun getName ()Ljava/lang/String; + public fun getParams ()Ljava/util/List; +} + +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 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 + public static synthetic fun encode$default (Lsk/ainet/models/bert/BertEncoderRuntime;[I[IILjava/lang/Object;)[F + public final fun exportTape (I)Lsk/ainet/lang/graph/DefaultExecutionTape; + public final fun forward ([I)Lsk/ainet/lang/tensor/Tensor; + public final fun getConfig ()Lsk/ainet/models/bert/BertModelConfig; + public final fun getDimensions ()I +} + +public final class sk/ainet/models/bert/BertExecutionMode : java/lang/Enum { + public static final field DIRECT Lsk/ainet/models/bert/BertExecutionMode; + public static final field OPTIMIZED Lsk/ainet/models/bert/BertExecutionMode; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/models/bert/BertExecutionMode; + public static fun values ()[Lsk/ainet/models/bert/BertExecutionMode; } public final class sk/ainet/models/bert/BertModelConfig { @@ -76,93 +69,6 @@ public final class sk/ainet/models/bert/BertNetworkLoader { public static final field INSTANCE Lsk/ainet/models/bert/BertNetworkLoader; } -public final class sk/ainet/models/bert/BertNetworkLoaderKt { - public static final fun buildBertTensorMap (Lsk/ainet/models/bert/BertRuntimeWeights;)Ljava/util/Map; -} - -public final class sk/ainet/models/bert/BertRuntime { - public fun (Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/bert/BertRuntimeWeights;Lkotlin/reflect/KClass;)V - public final fun encode ([I[I[I)Lsk/ainet/lang/tensor/Tensor; - public static synthetic fun encode$default (Lsk/ainet/models/bert/BertRuntime;[I[I[IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; - public final fun forward ([I[I)Lsk/ainet/lang/tensor/Tensor; - public static synthetic fun forward$default (Lsk/ainet/models/bert/BertRuntime;[I[IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; -} - -public final class sk/ainet/models/bert/BertRuntimeWeights { - public fun (Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V - public synthetic fun (Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;ILkotlin/jvm/internal/DefaultConstructorMarker;)V - public final fun component1 ()Lsk/ainet/models/bert/BertModelConfig; - public final fun component10 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component11 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component2 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component3 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component4 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component5 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component6 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component7 ()Ljava/util/List; - public final fun component8 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component9 ()Lsk/ainet/lang/tensor/Tensor; - public final fun copy (Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/models/bert/BertRuntimeWeights; - public static synthetic fun copy$default (Lsk/ainet/models/bert/BertRuntimeWeights;Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;ILjava/lang/Object;)Lsk/ainet/models/bert/BertRuntimeWeights; - public fun equals (Ljava/lang/Object;)Z - public final fun getConfig ()Lsk/ainet/models/bert/BertModelConfig; - public final fun getEmbeddingLayerNormBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getEmbeddingLayerNormWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getLayers ()Ljava/util/List; - public final fun getPoolerDenseBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getPoolerDenseWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getPositionEmbeddings ()Lsk/ainet/lang/tensor/Tensor; - public final fun getProjectionBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getProjectionWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getTokenTypeEmbeddings ()Lsk/ainet/lang/tensor/Tensor; - public final fun getWordEmbeddings ()Lsk/ainet/lang/tensor/Tensor; - public fun hashCode ()I - public fun toString ()Ljava/lang/String; -} - -public final class sk/ainet/models/bert/BertRuntimeWeightsKt { - public static final fun getMDBR_LEAF_IR_CONFIG ()Lsk/ainet/models/bert/BertModelConfig; -} - -public final class sk/ainet/models/bert/BertTensorNames { - public static final field EMBEDDING_LN_BIAS Ljava/lang/String; - public static final field EMBEDDING_LN_WEIGHT Ljava/lang/String; - public static final field INSTANCE Lsk/ainet/models/bert/BertTensorNames; - public static final field POOLER_DENSE_BIAS Ljava/lang/String; - public static final field POOLER_DENSE_WEIGHT Ljava/lang/String; - public static final field POSITION_EMBEDDINGS Ljava/lang/String; - public static final field PROJECTION_BIAS Ljava/lang/String; - public static final field PROJECTION_WEIGHT Ljava/lang/String; - public static final field TOKEN_TYPE_EMBEDDINGS Ljava/lang/String; - public static final field WORD_EMBEDDINGS Ljava/lang/String; - public final fun attnLayerNormBias (I)Ljava/lang/String; - public final fun attnLayerNormWeight (I)Ljava/lang/String; - public final fun attnOutputBias (I)Ljava/lang/String; - public final fun attnOutputWeight (I)Ljava/lang/String; - public final fun intermediateBias (I)Ljava/lang/String; - public final fun intermediateWeight (I)Ljava/lang/String; - public final fun keyBias (I)Ljava/lang/String; - public final fun keyWeight (I)Ljava/lang/String; - public final fun outputBias (I)Ljava/lang/String; - public final fun outputLayerNormBias (I)Ljava/lang/String; - public final fun outputLayerNormWeight (I)Ljava/lang/String; - public final fun outputWeight (I)Ljava/lang/String; - public final fun queryBias (I)Ljava/lang/String; - public final fun queryWeight (I)Ljava/lang/String; - public final fun valueBias (I)Ljava/lang/String; - public final fun valueWeight (I)Ljava/lang/String; -} - -public final class sk/ainet/models/bert/BertWeightLoaderKt { - public static final fun loadBertWeights (Ljava/util/List;Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/models/bert/BertModelConfig;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; - public static final fun loadBertWeights (Lsk/ainet/io/ParametersLoader;Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/models/bert/BertModelConfig;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; -} - -public final class sk/ainet/models/bert/BertWeightMapper { - public static final field INSTANCE Lsk/ainet/models/bert/BertWeightMapper; - public final fun map (Ljava/util/Map;Lsk/ainet/models/bert/BertModelConfig;)Lsk/ainet/models/bert/BertRuntimeWeights; -} - 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; @@ -194,4 +100,3 @@ public final class sk/ainet/models/bert/TokenizerOutput { public fun hashCode ()I public fun toString ()Ljava/lang/String; } - diff --git a/llm-inference/bert/api/jvm/bert.api b/llm-inference/bert/api/jvm/bert.api index fb332e3e..57b3bad3 100644 --- a/llm-inference/bert/api/jvm/bert.api +++ b/llm-inference/bert/api/jvm/bert.api @@ -1,47 +1,40 @@ -public final class sk/ainet/models/bert/BertIngestion { - public fun (Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/models/bert/BertModelConfig;)V - public final fun load (Lsk/ainet/io/ParametersLoader;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; +public final class sk/ainet/models/bert/BertConfigKt { + public static final fun getMDBR_LEAF_IR_CONFIG ()Lsk/ainet/models/bert/BertModelConfig; } -public final class sk/ainet/models/bert/BertLayerWeights { - public fun (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V - public final fun component1 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component10 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component11 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component12 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component13 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component14 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component15 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component16 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component2 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component3 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component4 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component5 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component6 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component7 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component8 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component9 ()Lsk/ainet/lang/tensor/Tensor; - public final fun copy (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/models/bert/BertLayerWeights; - public static synthetic fun copy$default (Lsk/ainet/models/bert/BertLayerWeights;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;ILjava/lang/Object;)Lsk/ainet/models/bert/BertLayerWeights; - public fun equals (Ljava/lang/Object;)Z - public final fun getAttnLayerNormBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getAttnLayerNormWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getAttnOutputBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getAttnOutputWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getIntermediateBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getIntermediateWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getKeyBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getKeyWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getOutputBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getOutputLayerNormBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getOutputLayerNormWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getOutputWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getQueryBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getQueryWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getValueBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getValueWeight ()Lsk/ainet/lang/tensor/Tensor; - public fun hashCode ()I - public fun toString ()Ljava/lang/String; +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 class sk/ainet/models/bert/BertEmbeddings : sk/ainet/lang/nn/Module, sk/ainet/lang/nn/topology/ModuleParameters { + public fun (Lsk/ainet/models/bert/BertModelConfig;Lkotlin/reflect/KClass;Ljava/lang/String;)V + public synthetic fun (Lsk/ainet/models/bert/BertModelConfig;Lkotlin/reflect/KClass;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun getModules ()Ljava/util/List; + public fun getName ()Ljava/lang/String; + public fun getParams ()Ljava/util/List; +} + +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 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 + public static synthetic fun encode$default (Lsk/ainet/models/bert/BertEncoderRuntime;[I[IILjava/lang/Object;)[F + public final fun exportTape (I)Lsk/ainet/lang/graph/DefaultExecutionTape; + public final fun forward ([I)Lsk/ainet/lang/tensor/Tensor; + public final fun getConfig ()Lsk/ainet/models/bert/BertModelConfig; + public final fun getDimensions ()I +} + +public final class sk/ainet/models/bert/BertExecutionMode : java/lang/Enum { + public static final field DIRECT Lsk/ainet/models/bert/BertExecutionMode; + public static final field OPTIMIZED Lsk/ainet/models/bert/BertExecutionMode; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/models/bert/BertExecutionMode; + public static fun values ()[Lsk/ainet/models/bert/BertExecutionMode; } public final class sk/ainet/models/bert/BertModelConfig { @@ -74,93 +67,10 @@ public final class sk/ainet/models/bert/BertModelConfig { public final class sk/ainet/models/bert/BertNetworkLoader { public static final field INSTANCE Lsk/ainet/models/bert/BertNetworkLoader; -} - -public final class sk/ainet/models/bert/BertNetworkLoaderKt { - public static final fun buildBertTensorMap (Lsk/ainet/models/bert/BertRuntimeWeights;)Ljava/util/Map; -} - -public final class sk/ainet/models/bert/BertRuntime { - public fun (Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/bert/BertRuntimeWeights;Lkotlin/reflect/KClass;)V - public final fun encode ([I[I[I)Lsk/ainet/lang/tensor/Tensor; - public static synthetic fun encode$default (Lsk/ainet/models/bert/BertRuntime;[I[I[IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; - public final fun forward ([I[I)Lsk/ainet/lang/tensor/Tensor; - public static synthetic fun forward$default (Lsk/ainet/models/bert/BertRuntime;[I[IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; -} - -public final class sk/ainet/models/bert/BertRuntimeWeights { - public fun (Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V - public synthetic fun (Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;ILkotlin/jvm/internal/DefaultConstructorMarker;)V - public final fun component1 ()Lsk/ainet/models/bert/BertModelConfig; - public final fun component10 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component11 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component2 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component3 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component4 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component5 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component6 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component7 ()Ljava/util/List; - public final fun component8 ()Lsk/ainet/lang/tensor/Tensor; - public final fun component9 ()Lsk/ainet/lang/tensor/Tensor; - public final fun copy (Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/models/bert/BertRuntimeWeights; - public static synthetic fun copy$default (Lsk/ainet/models/bert/BertRuntimeWeights;Lsk/ainet/models/bert/BertModelConfig;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;ILjava/lang/Object;)Lsk/ainet/models/bert/BertRuntimeWeights; - public fun equals (Ljava/lang/Object;)Z - public final fun getConfig ()Lsk/ainet/models/bert/BertModelConfig; - public final fun getEmbeddingLayerNormBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getEmbeddingLayerNormWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getLayers ()Ljava/util/List; - public final fun getPoolerDenseBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getPoolerDenseWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getPositionEmbeddings ()Lsk/ainet/lang/tensor/Tensor; - public final fun getProjectionBias ()Lsk/ainet/lang/tensor/Tensor; - public final fun getProjectionWeight ()Lsk/ainet/lang/tensor/Tensor; - public final fun getTokenTypeEmbeddings ()Lsk/ainet/lang/tensor/Tensor; - public final fun getWordEmbeddings ()Lsk/ainet/lang/tensor/Tensor; - public fun hashCode ()I - public fun toString ()Ljava/lang/String; -} - -public final class sk/ainet/models/bert/BertRuntimeWeightsKt { - public static final fun getMDBR_LEAF_IR_CONFIG ()Lsk/ainet/models/bert/BertModelConfig; -} - -public final class sk/ainet/models/bert/BertTensorNames { - public static final field EMBEDDING_LN_BIAS Ljava/lang/String; - public static final field EMBEDDING_LN_WEIGHT Ljava/lang/String; - public static final field INSTANCE Lsk/ainet/models/bert/BertTensorNames; - public static final field POOLER_DENSE_BIAS Ljava/lang/String; - public static final field POOLER_DENSE_WEIGHT Ljava/lang/String; - public static final field POSITION_EMBEDDINGS Ljava/lang/String; public static final field PROJECTION_BIAS Ljava/lang/String; public static final field PROJECTION_WEIGHT Ljava/lang/String; - public static final field TOKEN_TYPE_EMBEDDINGS Ljava/lang/String; - public static final field WORD_EMBEDDINGS Ljava/lang/String; - public final fun attnLayerNormBias (I)Ljava/lang/String; - public final fun attnLayerNormWeight (I)Ljava/lang/String; - public final fun attnOutputBias (I)Ljava/lang/String; - public final fun attnOutputWeight (I)Ljava/lang/String; - public final fun intermediateBias (I)Ljava/lang/String; - public final fun intermediateWeight (I)Ljava/lang/String; - public final fun keyBias (I)Ljava/lang/String; - public final fun keyWeight (I)Ljava/lang/String; - public final fun outputBias (I)Ljava/lang/String; - public final fun outputLayerNormBias (I)Ljava/lang/String; - public final fun outputLayerNormWeight (I)Ljava/lang/String; - public final fun outputWeight (I)Ljava/lang/String; - public final fun queryBias (I)Ljava/lang/String; - public final fun queryWeight (I)Ljava/lang/String; - public final fun valueBias (I)Ljava/lang/String; - public final fun valueWeight (I)Ljava/lang/String; -} - -public final class sk/ainet/models/bert/BertWeightLoaderKt { - public static final fun loadBertWeights (Ljava/util/List;Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/models/bert/BertModelConfig;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; - public static final fun loadBertWeights (Lsk/ainet/io/ParametersLoader;Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/models/bert/BertModelConfig;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; -} - -public final class sk/ainet/models/bert/BertWeightMapper { - public static final field INSTANCE Lsk/ainet/models/bert/BertWeightMapper; - public final fun map (Ljava/util/Map;Lsk/ainet/models/bert/BertModelConfig;)Lsk/ainet/models/bert/BertRuntimeWeights; + public final fun loadWeightTensors (Ljava/util/List;Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public final fun normalizeTensorNames (Ljava/util/List;)Ljava/util/List; } public final class sk/ainet/models/bert/HuggingFaceTokenizer : sk/ainet/apps/llm/Tokenizer { @@ -233,7 +143,7 @@ public final class sk/ainet/models/bert/java/KBertJava { } public final class sk/ainet/models/bert/java/KBertSession : java/lang/AutoCloseable { - public fun (Lsk/ainet/models/bert/BertRuntime;Lsk/ainet/models/bert/HuggingFaceTokenizer;)V + public fun (Lsk/ainet/models/bert/BertEncoderRuntime;Lsk/ainet/models/bert/HuggingFaceTokenizer;)V public fun close ()V public final fun encode (Ljava/lang/String;)[F public final fun similarity (Ljava/lang/String;Ljava/lang/String;)F diff --git a/llm-inference/bert/build.gradle.kts b/llm-inference/bert/build.gradle.kts index e5095eac..2284ca06 100644 --- a/llm-inference/bert/build.gradle.kts +++ b/llm-inference/bert/build.gradle.kts @@ -46,6 +46,10 @@ kotlin { implementation(libs.skainet.lang.core) implementation(libs.skainet.io.safetensors) implementation(libs.skainet.compile.core) + // Traced ComputeGraph execution for BertEncoderRuntime's OPTIMIZED + // mode (llm-core declares these as implementation — not transitive). + implementation(libs.skainet.compile.dag) + implementation(libs.skainet.compile.opt) implementation(libs.skainet.backend.cpu) implementation(libs.skainet.io.core) implementation(libs.kotlinx.coroutines) @@ -63,6 +67,8 @@ kotlin { implementation(libs.kotlinx.coroutines.test) implementation(libs.skainet.backend.cpu) implementation(libs.skainet.io.safetensors) + // StableHLO export smoke test (JVM-only, like gemma's MLIR dump test) + implementation(libs.skainet.compile.hlo) } } } 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 new file mode 100644 index 00000000..a130a16f --- /dev/null +++ b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertConfig.kt @@ -0,0 +1,70 @@ +package sk.ainet.models.bert + +/** + * Configuration for a BERT model architecture. + */ +public data class BertModelConfig( + val vocabSize: Int, + val hiddenSize: Int, + val numHiddenLayers: Int, + val numAttentionHeads: Int, + val intermediateSize: Int, + val maxPositionEmbeddings: Int = 512, + val typeVocabSize: Int = 2, + val layerNormEps: Double = 1e-12, + /** If non-null, a dense projection applied after pooling (sentence-transformers). */ + val projectionDim: Int? = null +) + +/** Default config for MongoDB/mdbr-leaf-ir (23M-param BERT embedding model). */ +public val MDBR_LEAF_IR_CONFIG: BertModelConfig = BertModelConfig( + vocabSize = 30522, + hiddenSize = 384, + numHiddenLayers = 6, + numAttentionHeads = 12, + intermediateSize = 1536, + maxPositionEmbeddings = 512, + typeVocabSize = 2, + layerNormEps = 1e-12, + projectionDim = 768 +) + +/** + * Parses a HuggingFace BERT `config.json` (and the optional sentence-transformers + * `2_Dense/config.json` projection head) into a [BertModelConfig]. + * + * Deliberately regex-based: the values are flat scalars, and this module stays + * free of serialization dependencies across its multiplatform targets. + */ +public object BertConfigParser { + + /** + * @param configJson content of the model's `config.json` + * @param denseConfigJson content of `2_Dense/config.json` when the snapshot + * ships a dense projection head; its `out_features` becomes + * [BertModelConfig.projectionDim] + */ + public fun parse(configJson: String, denseConfigJson: String? = null): BertModelConfig { + fun extractInt(source: String, key: String, default: Int): Int = + Regex("\"$key\"\\s*:\\s*(\\d+)").find(source)?.groupValues?.get(1)?.toIntOrNull() ?: default + + fun extractDouble(source: String, key: String, default: Double): Double = + Regex("\"$key\"\\s*:\\s*([\\d.eE\\-+]+)").find(source)?.groupValues?.get(1)?.toDoubleOrNull() ?: default + + val projectionDim = denseConfigJson + ?.let { extractInt(it, "out_features", 0) } + ?.takeIf { it > 0 } + + return BertModelConfig( + vocabSize = extractInt(configJson, "vocab_size", 30522), + hiddenSize = extractInt(configJson, "hidden_size", 768), + numHiddenLayers = extractInt(configJson, "num_hidden_layers", 12), + numAttentionHeads = extractInt(configJson, "num_attention_heads", 12), + intermediateSize = extractInt(configJson, "intermediate_size", 3072), + maxPositionEmbeddings = extractInt(configJson, "max_position_embeddings", 512), + typeVocabSize = extractInt(configJson, "type_vocab_size", 2), + layerNormEps = extractDouble(configJson, "layer_norm_eps", 1e-12), + projectionDim = projectionDim, + ) + } +} diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertEmbeddings.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertEmbeddings.kt new file mode 100644 index 00000000..7f105cac --- /dev/null +++ b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertEmbeddings.kt @@ -0,0 +1,118 @@ +package sk.ainet.models.bert + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.layers.Embedding +import sk.ainet.lang.nn.layers.EmbeddingAdapter +import sk.ainet.lang.nn.normalization.LayerNormalization +import sk.ainet.lang.nn.topology.ModuleParameter +import sk.ainet.lang.nn.topology.ModuleParameters +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.DType +import kotlin.reflect.KClass + +/** + * Complete BERT embeddings block: word + absolute position + token-type + * embeddings, summed and layer-normalized. + * + * `hidden = LayerNorm(wordEmb[ids] + posEmb[0..L-1] + typeEmb[0])` + * + * The three lookups are expressed with graph-safe ops only: + * - word embeddings via [Embedding]'s `ops.gather` (the input tensor carries + * the token ids), + * - position embeddings via `ops.narrow(table, 0, 0, L)` — rows `0..L-1` of + * the position table *are* the absolute-position vectors, so no index + * tensor is needed, + * - token-type embeddings as row 0 of the type table reshaped to `[dim]` and + * broadcast-added (the same `[L, dim] + [dim]` broadcast Linear uses for + * bias). + * + * Keeping position/type lookups index-free means a traced forward has exactly + * one non-parameter leaf (the token-id tensor), which makes compiled-mode + * input detection unambiguous and the exported graph a clean + * `tokens → hidden-states` encoder. + * + * Single-segment only: all positions use token-type row 0, which matches every + * sentence-embedding caller (they pass all-zero segment ids). Two-segment + * (cross-encoder) inputs are not supported by the DSL network. + * + * Weight-mapping paths (resolved by `BertSafeTensorsNameResolver`): + * - `…/embeddings/word_embeddings` → `bert.embeddings.word_embeddings.weight` + * - `…/embeddings` + `position_embeddings.weight` → `bert.embeddings.position_embeddings.weight` + * - `…/embeddings` + `token_type_embeddings.weight` → `bert.embeddings.token_type_embeddings.weight` + * - `…/embeddings/LayerNorm` → `bert.embeddings.LayerNorm.{weight,bias}` + */ +public class BertEmbeddings( + config: BertModelConfig, + dtype: KClass, + override val name: String = "embeddings", +) : Module(), ModuleParameters { + + private val dim = config.hiddenSize + + private val wordEmbeddings: EmbeddingAdapter = EmbeddingAdapter( + Embedding( + numEmbeddings = config.vocabSize, + embeddingDim = dim, + initWeight = voidTensor(dtype, config.vocabSize, dim), + name = "word_embeddings", + ) + ) + + private val layerNorm: LayerNormalization = LayerNormalization( + normalizedShape = intArrayOf(dim), + eps = config.layerNormEps, + name = "LayerNorm", + dtype = dtype, + ) + + override val params: List> = listOf( + ModuleParameter.WeightParameter( + "position_embeddings.weight", + voidTensor(dtype, config.maxPositionEmbeddings, dim), + ), + ModuleParameter.WeightParameter( + "token_type_embeddings.weight", + voidTensor(dtype, config.typeVocabSize, dim), + ), + ) + + override val modules: List> = listOf(wordEmbeddings, layerNorm) + + override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor { + val word = wordEmbeddings.forward(input, ctx) // gather → [L, dim] + val seqLen = word.shape[0] + + val positionTable = params[0].value + require(seqLen <= positionTable.shape[0]) { + "BertEmbeddings($name): sequence length $seqLen exceeds maxPositionEmbeddings ${positionTable.shape[0]}" + } + // Rows 0..L-1 of the position table are exactly the absolute-position + // embeddings for positions 0..L-1 — same result as gather([0..L-1]). + val pos = ctx.ops.narrow(positionTable, dim = 0, start = 0, length = seqLen) // [L, dim] + + // Segment 0 for every position: row 0 of the type table as a [dim] + // vector, broadcast over the sequence like a Linear bias. + val typeRow = ctx.ops.narrow(params[1].value, dim = 0, start = 0, length = 1) // [1, dim] + val type = ctx.ops.reshape(typeRow, Shape(dim)) // [dim] + + val summed = ctx.ops.add(ctx.ops.add(word, pos), type) + return layerNorm.forward(summed, ctx) + } + + private companion object { + @Suppress("UNCHECKED_CAST") + private fun voidTensor(dtype: KClass, rows: Int, cols: Int): Tensor = + VoidOpsTensor( + object : TensorData { + override val shape = Shape(rows, cols) + override fun get(vararg indices: Int): V = 0.0f as V + override fun set(vararg indices: Int, value: V) {} + }, + dtype, + ) + } +} 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 new file mode 100644 index 00000000..a3906f96 --- /dev/null +++ b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertEncoderRuntime.kt @@ -0,0 +1,342 @@ +package sk.ainet.models.bert + +import sk.ainet.apps.llm.getLLMOptimizationPipeline +import sk.ainet.apps.llm.graph.LLMFusedOpHandlers +import sk.ainet.apps.llm.weights.BertSafeTensorsNameResolver +import sk.ainet.context.ExecutionContext +import sk.ainet.io.weights.MappingConfig +import sk.ainet.io.weights.WeightMapper +import sk.ainet.io.weights.WeightNameResolver +import sk.ainet.io.weights.WeightTensor +import sk.ainet.lang.graph.DefaultExecutionTape +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.graph.exec.ComputeGraphExecutor +import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.topology.ModuleParameters +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +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.sqrt +import sk.ainet.lang.tensor.sum +import sk.ainet.lang.tensor.t +import sk.ainet.lang.tensor.times +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.Int32 +import kotlin.reflect.KClass + +/** Execution strategy for [BertEncoderRuntime]. */ +public enum class BertExecutionMode { + /** Run the module tree eagerly — the primary JVM path. */ + DIRECT, + + /** Trace the encoder into an optimized ComputeGraph and execute the graph. */ + OPTIMIZED, +} + +/** + * Encoder runtime for BERT sentence embeddings on the DSL path. + * + * Wraps a [bertNetwork] module (a complete `tokens → hidden-states` encoder) + * and adds what sentence embedding needs on top of the pure encoder graph: + * masked mean pooling, the optional sentence-transformers dense projection + * (`2_Dense`), and L2 normalization. Pooling and projection stay outside the + * DSL network on purpose — the traced/exported graph remains a clean encoder, + * and the pooling mask is dynamic per call. + * + * Intended use is one unpadded sequence per [encode] call; the attention mask + * only affects pooling (attention itself is bidirectional over the full + * sequence, exactly like the eager runtime this class replaces). + * + * Construction goes through [createBertEncoderRuntime], which maps checkpoint + * tensors into the module via [WeightMapper]. + */ +public class BertEncoderRuntime( + private val model: Module, + public val config: BertModelConfig, + private val ctx: ExecutionContext, + private val dtype: KClass, + private val projectionWeight: Tensor? = null, + private val projectionBias: Tensor? = null, + private val mode: BertExecutionMode = BertExecutionMode.DIRECT, +) { + + /** Output dimensionality of [encode]: projection out-features when present, else hidden size. */ + public val dimensions: Int get() = config.projectionDim ?: config.hiddenSize + + /** + * Full-sequence encoder forward: `[L]` token ids → `[L, hiddenSize]` + * hidden states. + */ + public fun forward(tokenIds: IntArray): Tensor = ctx.scratch.scope { + require(tokenIds.isNotEmpty()) { "BertEncoderRuntime: tokenIds must not be empty" } + when (mode) { + BertExecutionMode.DIRECT -> model.forward(tokenTensor(tokenIds), ctx) + BertExecutionMode.OPTIMIZED -> forwardOptimized(tokenIds) + } + } + + /** + * Encode tokens into a single embedding vector: forward → mean pooling + * (mask-weighted when [attentionMask] is given) → 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 + * @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" + } + 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) { + // sentence-transformers Dense head; bias is optional (LEAF models + // ship bias=false). The legacy eager runtime required both tensors + // and silently skipped the projection on bias-free heads — fixed here. + pooled = pooled.matmul(projectionWeight.t()) + if (projectionBias != null) pooled = pooled + projectionBias + } + + pooled = l2Normalize(pooled) + + val out = FloatArray(pooled.volume) + for (i in out.indices) out[i] = pooled.data[i] + return out + } + + // ------------------------------------------------------------------ + // OPTIMIZED mode: trace → ComputeGraph → fused execution. + // + // Traced graphs are shape-specialized: one compiled graph per sequence + // length, kept in a small LRU cache (re-tracing a 6×384 encoder is cheap, + // and embedding workloads cluster around few lengths after chunking). + // ------------------------------------------------------------------ + + private class CompiledEncoder( + val executor: ComputeGraphExecutor, + val weightMap: Map>, + val inputNodeId: String, + ) + + private val compiledBySeqLen = LinkedHashMap() + + /** + * Pre-warm the OPTIMIZED cache for [sampleSeqLen] and return compile + * diagnostics. Optional — [forward] compiles lazily per sequence length. + */ + public fun compile(sampleSeqLen: Int = 8): List { + val diagnostics = mutableListOf() + compiledFor(IntArray(sampleSeqLen) { it % config.vocabSize }, diagnostics) + return diagnostics + } + + /** + * Trace one encoder forward at [seqLen] and return the recorded tape — + * the entry point for graph export (StableHLO and friends). + */ + public fun exportTape(seqLen: Int): DefaultExecutionTape { + val tapingCtx = DefaultGraphExecutionContext.tape(baseOps = ctx.ops) + tapingCtx.startRecording() + model.forward(tokenTensor(IntArray(seqLen) { it % config.vocabSize }), tapingCtx) + val tape = tapingCtx.stopRecording() + return tape as? DefaultExecutionTape + ?: error("Expected DefaultExecutionTape but got ${tape?.let { it::class.simpleName }}") + } + + private fun forwardOptimized(tokenIds: IntArray): Tensor { + val compiled = compiledBySeqLen.getOrPut(tokenIds.size) { + if (compiledBySeqLen.size >= COMPILED_CACHE_LIMIT) { + compiledBySeqLen.remove(compiledBySeqLen.keys.first()) + } + compiledFor(tokenIds, mutableListOf()) + } + + // Feed the Int32 index tensor the traced gather consumes, plus every + // weight leaf resolved at compile time. + val input = ctx.fromIntArray(Shape(tokenIds.size), Int32::class, tokenIds) + val inputMap = mutableMapOf>() + @Suppress("UNCHECKED_CAST") + inputMap[compiled.inputNodeId] = input as Tensor + for ((nodeId, tensor) in compiled.weightMap) { + @Suppress("UNCHECKED_CAST") + inputMap[nodeId] = tensor as Tensor + } + + val results = compiled.executor.execute(inputMap) + @Suppress("UNCHECKED_CAST") + return results.values.lastOrNull() as? Tensor + ?: error("BertEncoderRuntime: graph execution produced no outputs") + } + + private fun compiledFor(sampleTokenIds: IntArray, diagnostics: MutableList): CompiledEncoder { + // Correct replay handlers (incl. the axes-aware permute override the + // multi-token attention trace depends on). + LLMFusedOpHandlers.registerAll() + + // 1. Record one eager forward through a tape-recording context. + val tapingCtx = DefaultGraphExecutionContext.tape(baseOps = ctx.ops) + tapingCtx.startRecording() + model.forward(tokenTensor(sampleTokenIds), tapingCtx) + val tape = tapingCtx.stopRecording() + val tapeObj = tape as? DefaultExecutionTape + ?: error("Expected DefaultExecutionTape but got ${tape?.let { it::class.simpleName }}") + val session = tapeObj.session + + // 2. Tape → raw ComputeGraph (weights stay placeholder inputs). + val rawGraph = tapeObj.toComputeGraph( + synthesizeExternalInputs = true, + inputTensorIds = emptySet(), + embedConstants = false, + ) + diagnostics.add("Traced graph: ${rawGraph.nodes.size} nodes, ${rawGraph.edges.size} edges") + + // 3. Optimize unless the raw graph fails validation. + val graph = try { + rawGraph.getTopologicalOrder() + val result = getLLMOptimizationPipeline().optimize(rawGraph) + diagnostics.addAll(result.passResults.flatMap { it.diagnostics }) + result.graph + } catch (e: IllegalStateException) { + diagnostics.add("WARNING: raw traced graph failed validation (${e.message}) — executing unoptimized") + rawGraph + } + + // 4. Split graph leaves into the token input vs static tensors. + // Module parameters are known by identity. Non-parameter leaves are + // either trace-time constants the modules created during forward + // (e.g. attention scale/mask tensors — valid for this seqLen, so + // they ride along in the static map) or the token input. The token + // input is the unique Int32 leaf: the index tensor Embedding + // creates for gather. (BertEmbeddings keeps position/type lookups + // index-free precisely so this stays unambiguous.) + val paramTensors = HashSet>() + collectParams(model, paramTensors) + + val weightMap = mutableMapOf>() + val inputCandidates = mutableListOf() + for (node in graph.nodes) { + if (node.operationName !in setOf("input", "weight", "parameter", "constant")) continue + val outputSpec = node.outputs.firstOrNull() ?: continue + val tensor = session.resolve(outputSpec.name) ?: continue + when { + paramTensors.any { it === tensor } -> weightMap[node.id] = tensor + tensor.dtype == Int32::class -> inputCandidates += node.id + else -> weightMap[node.id] = tensor + } + } + require(inputCandidates.size == 1) { + "BertEncoderRuntime: expected exactly one Int32 graph leaf (the token input), " + + "found ${inputCandidates.size}: $inputCandidates" + } + diagnostics.add("Input node: ${inputCandidates.first()}, weight map: ${weightMap.size} entries") + + return CompiledEncoder( + executor = ComputeGraphExecutor(graph, ctx.ops), + weightMap = weightMap, + inputNodeId = inputCandidates.first(), + ) + } + + private fun collectParams(module: Module<*, *>, out: MutableSet>) { + if (module is ModuleParameters<*, *>) { + module.params.forEach { out.add(it.value) } + } + module.modules.forEach { collectParams(it, out) } + } + + private companion object { + private const val COMPILED_CACHE_LIMIT = 16 + } + + private fun tokenTensor(tokenIds: IntArray): Tensor = + ctx.fromFloatArray( + Shape(tokenIds.size), dtype, + FloatArray(tokenIds.size) { tokenIds[it].toFloat() }, + ) + + private fun l2Normalize(tensor: Tensor): Tensor { + val squared = tensor * tensor + val sumSquared = squared.sum() + val norm = (sumSquared + 1e-12).sqrt() + return tensor / norm + } +} + +/** + * Build a [BertEncoderRuntime] from checkpoint tensors: constructs + * `bertNetwork(config)`, maps every DSL parameter by name via [resolver] + * (strict — shape fallback stays off so same-shaped Q/K/V tensors can't + * cross-wire), and pulls the optional `2_Dense` projection pair + * (`linear.weight` / `linear.bias`) out of [tensors] for the runtime. + * + * @param tensors as produced by [BertNetworkLoader.loadWeightTensors] + */ +public inline fun createBertEncoderRuntime( + config: BertModelConfig, + tensors: List>, + ctx: ExecutionContext, + resolver: WeightNameResolver = BertSafeTensorsNameResolver(), + mode: BertExecutionMode = BertExecutionMode.DIRECT, + debug: Boolean = false, +): BertEncoderRuntime { + val model = bertNetwork(config) + + val result = WeightMapper.applyWeights( + model, tensors, + MappingConfig( + usePathBasedMatching = false, + fallbackToShapeMatching = false, + debug = debug, + nameResolver = resolver, + ) + ) + require(result.mapped == result.total) { + buildString { + appendLine("Failed to map ${result.total - result.mapped}/${result.total} BERT parameters:") + result.missingParams.forEach { appendLine(" - $it") } + if (result.unusedTensors.isNotEmpty()) { + appendLine("Unused tensors (${result.unusedTensors.size}):") + result.unusedTensors.take(10).forEach { appendLine(" - $it") } + } + }.trim() + } + + val projectionWeight = tensors.firstOrNull { it.name == BertNetworkLoader.PROJECTION_WEIGHT }?.tensor + val projectionBias = tensors.firstOrNull { it.name == BertNetworkLoader.PROJECTION_BIAS }?.tensor + if (config.projectionDim != null) { + requireNotNull(projectionWeight) { + "config.projectionDim=${config.projectionDim} but no ${BertNetworkLoader.PROJECTION_WEIGHT} tensor was provided" + } + } + require(projectionBias == null || projectionWeight != null) { + "${BertNetworkLoader.PROJECTION_BIAS} provided without ${BertNetworkLoader.PROJECTION_WEIGHT}" + } + + return BertEncoderRuntime( + model = model, + config = config, + ctx = ctx, + dtype = T::class, + projectionWeight = projectionWeight, + projectionBias = projectionBias, + mode = mode, + ) +} diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertIngestion.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertIngestion.kt deleted file mode 100644 index fa35419b..00000000 --- a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertIngestion.kt +++ /dev/null @@ -1,27 +0,0 @@ -package sk.ainet.models.bert - -import sk.ainet.context.ExecutionContext -import sk.ainet.io.ParametersLoader -import sk.ainet.lang.types.DType -import kotlin.reflect.KClass - -/** - * Facade for loading BERT models from any format via the generic [ParametersLoader] interface. - * - * Follows the [LlamaIngestion] pattern but is IO-format agnostic — the caller provides - * a [ParametersLoader] which can be backed by SafeTensors, GGUF, ONNX, etc. - */ -public class BertIngestion( - private val ctx: ExecutionContext, - private val dtype: KClass, - private val config: BertModelConfig -) { - /** - * Load BERT weights from a generic [ParametersLoader]. - * Works with any format (SafeTensors, GGUF, ONNX) as long as the tensor names - * follow HuggingFace conventions. - */ - public suspend fun load(loader: ParametersLoader): BertRuntimeWeights { - return loadBertWeights(loader, ctx, dtype, config) - } -} diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertNetworkDef.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertNetworkDef.kt index 4ecd397c..d05d3fa5 100644 --- a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertNetworkDef.kt +++ b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertNetworkDef.kt @@ -5,7 +5,6 @@ import sk.ainet.lang.nn.DefaultNeuralNetworkExecutionContext import sk.ainet.lang.nn.Module import sk.ainet.lang.nn.dsl.NeuralNetworkDslImpl import sk.ainet.lang.nn.dsl.StageImpl -import sk.ainet.lang.nn.dsl.embedding import sk.ainet.lang.nn.dsl.multiHeadAttention import sk.ainet.lang.nn.dsl.residual import sk.ainet.lang.nn.dsl.sequential @@ -15,15 +14,14 @@ import sk.ainet.lang.types.DType /** * BERT architecture defined via the network DSL. * - * Replaces the hand-coded [BertRuntime] with a declarative definition. - * - * Architecture: Embedding → LayerNorm → + * Architecture: BertEmbeddings (word + position + token_type + LayerNorm) → * N × (MHA(bidirectional, bias) → Residual → LayerNorm → * Dense → GeLU → Dense → Residual → LayerNorm) * - * Note: BERT has 3 embedding types (word, position, token_type) but only - * word embeddings are defined here. Position and token_type are handled - * by the runtime's forward logic as additive lookups. + * The returned module is a complete `tokens → hidden-states` encoder: feed it + * a `[L]`-shaped token-id tensor and it produces `[L, hiddenSize]` hidden + * states. Pooling and the optional sentence-transformers projection live in + * [BertEncoderRuntime], keeping the traced graph a pure encoder. */ public inline fun bertNetwork( config: BertModelConfig @@ -32,41 +30,47 @@ public inline fun bertNetwork( val nHeads = config.numAttentionHeads val nLayers = config.numHiddenLayers val ffnDim = config.intermediateSize - val vocabSize = config.vocabSize val eps = config.layerNormEps return sequential { val dslImpl = this as NeuralNetworkDslImpl val nnCtx = DefaultNeuralNetworkExecutionContext() - // Embedding stage: word embeddings + LayerNorm - val embStage = StageImpl(nnCtx, "embeddings", T::class) - embStage.embedding(vocabSize, dim, id = "word_embeddings") - embStage.layerNorm(intArrayOf(dim), eps, id = "LayerNorm") - dslImpl.modules += HybridTransformerBlock(embStage.modules.toList(), name = "embeddings") + // Complete embeddings block: word + position + token_type, LayerNorm + dslImpl.modules += BertEmbeddings(config, T::class) - // Encoder layers — use TransformerBlock for residual connections + // Encoder layers. BERT is POST-norm: each residual adds the value the + // sub-block started from, and LayerNorm sits AFTER the add: + // h1 = LN(x + MHA(x)); h2 = LN(h1 + FFN(h1)) + // The transformer blocks wire each ResidualAdd to the value at the + // start of its residual segment (the input of the module right after + // the previous ResidualAdd) — correct for pre-norm decoder stacks, but + // for BERT the FFN residual must start AT the post-attention LayerNorm + // output. Splitting the layer into two blocks puts each residual + // boundary exactly there: within a block, the first segment starts at + // the block input. for (layer in 0 until nLayers) { - val stage = StageImpl(nnCtx, "encoder.layer.$layer", T::class) + val attnStage = StageImpl(nnCtx, "encoder.layer.$layer.attn", T::class) // Self-attention (bidirectional, no causal mask, with bias) - stage.multiHeadAttention( + attnStage.multiHeadAttention( dim = dim, nHeads = nHeads, causal = false, bias = true, id = "attention" ) - stage.residual() - stage.layerNorm(intArrayOf(dim), eps, id = "attn_ln") - - // GeLU FFN - stage.dense(ffnDim, id = "intermediate") - stage.activation { it.gelu() } - stage.dense(dim, id = "output") - stage.residual() - stage.layerNorm(intArrayOf(dim), eps, id = "output_ln") + attnStage.residual() + attnStage.layerNorm(intArrayOf(dim), eps, id = "attn_ln") + dslImpl.modules += HybridTransformerBlock(attnStage.modules.toList(), name = "encoder.layer.$layer.attn") - dslImpl.modules += HybridTransformerBlock(stage.modules.toList(), name = "encoder.layer.$layer") + // GeLU FFN; residual adds this block's input = post-attention LN output + val ffnStage = StageImpl(nnCtx, "encoder.layer.$layer.ffn", T::class) + ffnStage.dense(ffnDim, id = "intermediate") + ffnStage.activation { it.gelu() } + ffnStage.dense(dim, id = "output") + ffnStage.residual() + ffnStage.layerNorm(intArrayOf(dim), eps, id = "output_ln") + dslImpl.modules += HybridTransformerBlock(ffnStage.modules.toList(), name = "encoder.layer.$layer.ffn") } } } diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertNetworkLoader.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertNetworkLoader.kt index dd0c81ed..ffbf4bdc 100644 --- a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertNetworkLoader.kt +++ b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertNetworkLoader.kt @@ -1,8 +1,8 @@ package sk.ainet.models.bert +import sk.ainet.apps.llm.weights.BertSafeTensorsNameResolver import sk.ainet.context.ExecutionContext import sk.ainet.io.ParametersLoader -import sk.ainet.io.weights.BertSafeTensorsNameResolver import sk.ainet.io.weights.MappingConfig import sk.ainet.io.weights.WeightMapper import sk.ainet.io.weights.WeightTensor @@ -13,88 +13,74 @@ import kotlin.reflect.KClass /** * End-to-end loader that builds a `bertNetwork()` module and populates it - * with weights from SafeTensors (HuggingFace) format via [WeightMapper] + [BertSafeTensorsNameResolver]. + * with weights from SafeTensors (HuggingFace) format via [WeightMapper] + + * [BertSafeTensorsNameResolver]. * * Usage: * ```kotlin - * val model = BertNetworkLoader.fromTensorMap(config, tensors) - * - * // Or from existing BertRuntimeWeights - * val model = BertNetworkLoader.fromRuntimeWeights(config, runtimeWeights) + * val tensors = BertNetworkLoader.loadWeightTensors(loaders, ctx, FP32::class) + * val model = BertNetworkLoader.fromWeightTensors(config, tensors) * ``` - * - * Note: BERT's position and token_type embeddings are not part of the DSL network - * (they require additive lookups during forward pass). Only word embeddings, - * LayerNorm, attention, and FFN weights are mapped. */ -@PublishedApi -internal fun buildBertTensorMap( - weights: BertRuntimeWeights -): Map> { - val map = mutableMapOf>() +public object BertNetworkLoader { - map[BertTensorNames.WORD_EMBEDDINGS] = weights.wordEmbeddings - map[BertTensorNames.EMBEDDING_LN_WEIGHT] = weights.embeddingLayerNormWeight - map[BertTensorNames.EMBEDDING_LN_BIAS] = weights.embeddingLayerNormBias + /** The sentence-transformers dense-projection tensor names (2_Dense/model.safetensors). */ + public const val PROJECTION_WEIGHT: String = "linear.weight" + public const val PROJECTION_BIAS: String = "linear.bias" - weights.layers.forEachIndexed { i, layer -> - map[BertTensorNames.queryWeight(i)] = layer.queryWeight - map[BertTensorNames.queryBias(i)] = layer.queryBias - map[BertTensorNames.keyWeight(i)] = layer.keyWeight - map[BertTensorNames.keyBias(i)] = layer.keyBias - map[BertTensorNames.valueWeight(i)] = layer.valueWeight - map[BertTensorNames.valueBias(i)] = layer.valueBias - map[BertTensorNames.attnOutputWeight(i)] = layer.attnOutputWeight - map[BertTensorNames.attnOutputBias(i)] = layer.attnOutputBias - map[BertTensorNames.attnLayerNormWeight(i)] = layer.attnLayerNormWeight - map[BertTensorNames.attnLayerNormBias(i)] = layer.attnLayerNormBias - map[BertTensorNames.intermediateWeight(i)] = layer.intermediateWeight - map[BertTensorNames.intermediateBias(i)] = layer.intermediateBias - map[BertTensorNames.outputWeight(i)] = layer.outputWeight - map[BertTensorNames.outputBias(i)] = layer.outputBias - map[BertTensorNames.outputLayerNormWeight(i)] = layer.outputLayerNormWeight - map[BertTensorNames.outputLayerNormBias(i)] = layer.outputLayerNormBias + /** + * Collect all tensors from [loaders] (base checkpoint plus the optional + * `2_Dense/model.safetensors` projection file) into a flat [WeightTensor] + * list with normalized HF names. + * + * Name normalization: plain BERT checkpoints prefix every tensor with + * `bert.`; sentence-transformers checkpoints (e.g. MongoDB/mdbr-leaf) ship + * bare names. When no key carries the prefix, every tensor except the + * projection pair gains it, so [BertSafeTensorsNameResolver]'s + * `bert.`-prefixed output matches either layout. + */ + public suspend fun loadWeightTensors( + loaders: List, + ctx: ExecutionContext, + dtype: KClass, + ): List> { + val collected = mutableListOf>() + loaders.forEach { loader -> + loader.load(ctx, dtype) { name, tensor -> + collected += WeightTensor( + name = name, + shape = tensor.shape.dimensions.toList(), + tensor = tensor, + ) + } + } + return normalizeTensorNames(collected) } - return map -} - -public object BertNetworkLoader { - /** - * Build a BERT network from a flat tensor map with HuggingFace tensor names. - * - * @param config BERT model configuration - * @param tensors Map of HF tensor names to tensors (e.g. from SafeTensors/ONNX loading) - * @param debug Whether to print debug mapping information + * Build a BERT network and map [tensors] (as produced by + * [loadWeightTensors]) into it. Position/token_type/pooler tensors that + * the checkpoint carries beyond the DSL parameters are reported as unused, + * which is expected; every DSL parameter must be mapped. */ - public inline fun fromTensorMap( + public inline fun fromWeightTensors( config: BertModelConfig, - tensors: Map>, - debug: Boolean = false + tensors: List>, + debug: Boolean = false, ): Module { val model = bertNetwork(config) - val weightTensors = tensors.map { (name, tensor) -> - WeightTensor( - name = name, - shape = tensor.shape.dimensions.toList(), - tensor = tensor - ) - } - val mappingConfig = MappingConfig( + // Path-based ONNX-style matching never applies to HF BERT names; + // shape fallback stays off so same-shaped Q/K/V can't cross-wire. usePathBasedMatching = false, fallbackToShapeMatching = false, debug = debug, nameResolver = BertSafeTensorsNameResolver() ) - val result = WeightMapper.applyWeights(model, weightTensors, mappingConfig) + val result = WeightMapper.applyWeights(model, tensors, mappingConfig) - // BERT tensors include position_embeddings, token_type_embeddings, pooler, etc. - // that are not in the DSL model — these will appear as unused tensors, which is fine. - // But all DSL parameters must be mapped. require(result.mapped == result.total) { buildString { appendLine("Failed to map ${result.total - result.mapped}/${result.total} BERT parameters:") @@ -110,18 +96,35 @@ public object BertNetworkLoader { } /** - * Build a BERT network from existing [BertRuntimeWeights]. - * - * Converts the structured weights back to a flat HF-named tensor map - * and feeds it through the standard mapping pipeline. + * Build a BERT network from a flat tensor map with HuggingFace tensor names. */ - public inline fun fromRuntimeWeights( - weights: BertRuntimeWeights, + public inline fun fromTensorMap( + config: BertModelConfig, + tensors: Map>, debug: Boolean = false ): Module { - @Suppress("UNCHECKED_CAST") - val tensors = buildBertTensorMap(weights) as Map> + val weightTensors = normalizeTensorNames( + tensors.map { (name, tensor) -> + WeightTensor( + name = name, + shape = tensor.shape.dimensions.toList(), + tensor = tensor + ) + } + ) + return fromWeightTensors(config, weightTensors, debug) + } - return fromTensorMap(weights.config, tensors, debug) + /** See [loadWeightTensors] for the normalization contract. */ + public fun normalizeTensorNames( + tensors: List>, + ): List> { + if (tensors.any { it.name.startsWith("bert.") }) return tensors + return tensors.map { wt -> + when (wt.name) { + PROJECTION_WEIGHT, PROJECTION_BIAS -> wt + else -> wt.copy(name = "bert.${wt.name}") + } + } } } diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertRuntime.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertRuntime.kt deleted file mode 100644 index efc271d9..00000000 --- a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertRuntime.kt +++ /dev/null @@ -1,247 +0,0 @@ -package sk.ainet.models.bert - -import sk.ainet.context.ExecutionContext -import sk.ainet.lang.nn.layers.Embedding -import sk.ainet.lang.nn.normalization.LayerNormalization -import sk.ainet.lang.tensor.Shape -import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.tensor.div -import sk.ainet.lang.tensor.gelu -import sk.ainet.lang.tensor.matmul -import sk.ainet.lang.tensor.mean -import sk.ainet.lang.tensor.minus -import sk.ainet.lang.tensor.narrow -import sk.ainet.lang.tensor.plus -import sk.ainet.lang.tensor.reshape -import sk.ainet.lang.tensor.softmax -import sk.ainet.lang.tensor.sqrt -import sk.ainet.lang.tensor.sum -import sk.ainet.lang.tensor.t -import sk.ainet.lang.tensor.times -import sk.ainet.lang.tensor.unsqueeze -import sk.ainet.lang.tensor.data.FloatArrayTensorData -import sk.ainet.lang.types.DType -import sk.ainet.lang.types.FP32 -import kotlin.math.sqrt as mathSqrt -import kotlin.reflect.KClass - -/** - * BERT encoder runtime. Produces contextual embeddings from token IDs. - * - * Follows the [LlamaRuntime] pattern: direct tensor ops, no Module composition for encoder layers. - * Uses [Embedding] for lookup tables and [LayerNormalization] for norms. - */ -@Deprecated( - message = "Use OptimizedLLMRuntime with bertNetwork() instead. " + - "See docs/optimizable-LLM-NNs-DAG.md for migration guide.", - replaceWith = ReplaceWith( - "OptimizedLLMRuntime.create(bertNetwork(config), tensors, resolver, ctx)", - "sk.ainet.apps.llm.OptimizedLLMRuntime" - ) -) -public class BertRuntime( - private val ctx: ExecutionContext, - private val weights: BertRuntimeWeights, - private val dtype: KClass -) { - private val config: BertModelConfig get() = weights.config - private val headDim: Int = config.hiddenSize / config.numAttentionHeads - - init { - require(headDim * config.numAttentionHeads == config.hiddenSize) { - "hiddenSize (${config.hiddenSize}) must be divisible by numAttentionHeads (${config.numAttentionHeads})" - } - } - - // Embedding modules - private val wordEmbedding = Embedding( - numEmbeddings = config.vocabSize, - embeddingDim = config.hiddenSize, - initWeight = weights.wordEmbeddings, - name = "word_embeddings" - ) - private val positionEmbedding = Embedding( - numEmbeddings = config.maxPositionEmbeddings, - embeddingDim = config.hiddenSize, - initWeight = weights.positionEmbeddings, - name = "position_embeddings" - ) - private val tokenTypeEmbedding = Embedding( - numEmbeddings = config.typeVocabSize, - embeddingDim = config.hiddenSize, - initWeight = weights.tokenTypeEmbeddings, - name = "token_type_embeddings" - ) - private val embeddingLayerNorm = LayerNormalization( - normalizedShape = intArrayOf(config.hiddenSize), - eps = config.layerNormEps, - elementwiseAffine = true, - name = "embeddings.LayerNorm", - initGamma = weights.embeddingLayerNormWeight, - initBeta = weights.embeddingLayerNormBias - ) - - // Per-layer LayerNorm modules - private val attnLayerNorms = weights.layers.mapIndexed { i, layer -> - LayerNormalization( - normalizedShape = intArrayOf(config.hiddenSize), - eps = config.layerNormEps, - elementwiseAffine = true, - name = "encoder.layer.$i.attention.output.LayerNorm", - initGamma = layer.attnLayerNormWeight, - initBeta = layer.attnLayerNormBias - ) - } - private val outputLayerNorms = weights.layers.mapIndexed { i, layer -> - LayerNormalization( - normalizedShape = intArrayOf(config.hiddenSize), - eps = config.layerNormEps, - elementwiseAffine = true, - name = "encoder.layer.$i.output.LayerNorm", - initGamma = layer.outputLayerNormWeight, - initBeta = layer.outputLayerNormBias - ) - } - - /** - * Full encoder forward pass: embeddings → encoder layers → hidden states. - * - * @param tokenIds token IDs, shape [seqLen] - * @param tokenTypeIds optional segment IDs, shape [seqLen] (defaults to all zeros) - * @return hidden states tensor of shape [seqLen, hiddenSize] - */ - public fun forward(tokenIds: IntArray, tokenTypeIds: IntArray? = null): Tensor = - ctx.scratch.scope { - // ScratchPool scope for the whole forward pass: upstream SIMD - // kernels (matmul, dequant) acquire their per-call workspace - // from ctx.scratch and the buffers are returned to the pool on - // scope exit. With the default NoopScratchPool this is a plain - // pass-through; with a SizeClassedScratchPool it eliminates per- - // forward FloatArray allocations on the embedding hot path. - val seqLen = tokenIds.size - val typeIds = tokenTypeIds ?: IntArray(seqLen) { 0 } - val positionIds = IntArray(seqLen) { it } - - // Embedding: word + position + token_type - val wordEmb = wordEmbedding.forward(tokenIds, ctx) - val posEmb = positionEmbedding.forward(positionIds, ctx) - val typeEmb = tokenTypeEmbedding.forward(typeIds, ctx) - - var hidden = wordEmb + posEmb + typeEmb - hidden = embeddingLayerNorm.forward(hidden, ctx) - - // Encoder layers - for (i in weights.layers.indices) { - hidden = runEncoderLayer(i, hidden) - } - - hidden - } - - /** - * Encode text tokens into a single embedding vector (mean pooling + optional projection + L2 norm). - * - * @param tokenIds token IDs including [CLS] and [SEP] - * @param attentionMask 1 for real tokens, 0 for padding. If null, all tokens attend. - * @param tokenTypeIds optional segment IDs - * @return normalized embedding vector of shape [projDim] or [hiddenSize] - */ - public fun encode( - tokenIds: IntArray, - attentionMask: IntArray? = null, - tokenTypeIds: IntArray? = null - ): Tensor { - val hiddenStates = forward(tokenIds, tokenTypeIds) - val seqLen = tokenIds.size - - // Mean pooling with attention mask - var pooled = if (attentionMask != null) { - 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) - } - - // Optional dense projection (sentence-transformers 2_Dense) - val projW = weights.projectionWeight - val projB = weights.projectionBias - if (projW != null && projB != null) { - pooled = pooled.matmul(projW.t()) + projB - } - - // L2 normalize - pooled = l2Normalize(pooled) - - return pooled - } - - private fun runEncoderLayer(layerIdx: Int, input: Tensor): Tensor { - val layer = weights.layers[layerIdx] - - // Self-attention: Q, K, V projections - val q = input.matmul(layer.queryWeight.t()) + layer.queryBias - val k = input.matmul(layer.keyWeight.t()) + layer.keyBias - val v = input.matmul(layer.valueWeight.t()) + layer.valueBias - - // Multi-head attention - val attnOutput = multiHeadAttention(q, k, v) - - // Attention output projection + residual + LayerNorm - val attnProj = attnOutput.matmul(layer.attnOutputWeight.t()) + layer.attnOutputBias - val attnResidual = input + attnProj - val attnNormed = attnLayerNorms[layerIdx].forward(attnResidual, ctx) - - // FFN: intermediate (dense + GELU) then output (dense) - val intermediate = (attnNormed.matmul(layer.intermediateWeight.t()) + layer.intermediateBias).gelu() - val ffnOutput = intermediate.matmul(layer.outputWeight.t()) + layer.outputBias - - // FFN residual + LayerNorm - val ffnResidual = attnNormed + ffnOutput - return outputLayerNorms[layerIdx].forward(ffnResidual, ctx) - } - - /** - * Multi-head attention using narrow/concat for head splitting. - * Input Q,K,V are [seqLen, hiddenSize], output is [seqLen, hiddenSize]. - * No causal mask (BERT is bidirectional). - */ - private fun multiHeadAttention( - q: Tensor, - k: Tensor, - v: Tensor - ): Tensor { - val numHeads = config.numAttentionHeads - val scale = mathSqrt(headDim.toDouble()).toFloat() - - val headOutputs = ArrayList>(numHeads) - for (h in 0 until numHeads) { - val offset = h * headDim - // Extract head slice: [seqLen, headDim] - val qh = q.narrow(dim = 1, start = offset, length = headDim) - val kh = k.narrow(dim = 1, start = offset, length = headDim) - val vh = v.narrow(dim = 1, start = offset, length = headDim) - - // Scaled dot-product attention: softmax(Q @ K^T / sqrt(d)) @ V - val scores = qh.matmul(kh.t()) / scale // [seqLen, seqLen] - val attnWeights = scores.softmax(dim = 1) // softmax over key dimension - val headOut = attnWeights.matmul(vh) // [seqLen, headDim] - headOutputs.add(headOut) - } - - // Concatenate heads along feature dimension - return q.ops.concat(headOutputs, dim = 1) // [seqLen, hiddenSize] - } - - private fun l2Normalize(tensor: Tensor): Tensor { - val squared = tensor * tensor - val sumSquared = squared.sum() - val norm = (sumSquared + 1e-12).sqrt() - return tensor / norm - } -} diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertRuntimeWeights.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertRuntimeWeights.kt deleted file mode 100644 index c1e987ec..00000000 --- a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertRuntimeWeights.kt +++ /dev/null @@ -1,79 +0,0 @@ -package sk.ainet.models.bert - -import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.types.DType - -/** - * Configuration for a BERT model architecture. - */ -public data class BertModelConfig( - val vocabSize: Int, - val hiddenSize: Int, - val numHiddenLayers: Int, - val numAttentionHeads: Int, - val intermediateSize: Int, - val maxPositionEmbeddings: Int = 512, - val typeVocabSize: Int = 2, - val layerNormEps: Double = 1e-12, - /** If non-null, a dense projection applied after pooling (sentence-transformers). */ - val projectionDim: Int? = null -) - -/** Default config for MongoDB/mdbr-leaf-ir (23M-param BERT embedding model). */ -public val MDBR_LEAF_IR_CONFIG: BertModelConfig = BertModelConfig( - vocabSize = 30522, - hiddenSize = 384, - numHiddenLayers = 6, - numAttentionHeads = 12, - intermediateSize = 1536, - maxPositionEmbeddings = 512, - typeVocabSize = 2, - layerNormEps = 1e-12, - projectionDim = 768 -) - -/** - * Per-layer weights for a BERT encoder layer. - */ -public data class BertLayerWeights( - // Self-attention - val queryWeight: Tensor, - val queryBias: Tensor, - val keyWeight: Tensor, - val keyBias: Tensor, - val valueWeight: Tensor, - val valueBias: Tensor, - // Attention output - val attnOutputWeight: Tensor, - val attnOutputBias: Tensor, - val attnLayerNormWeight: Tensor, - val attnLayerNormBias: Tensor, - // FFN (intermediate + output) - val intermediateWeight: Tensor, - val intermediateBias: Tensor, - val outputWeight: Tensor, - val outputBias: Tensor, - val outputLayerNormWeight: Tensor, - val outputLayerNormBias: Tensor -) - -/** - * Complete BERT runtime weights: embeddings + encoder layers + optional pooler/projection. - */ -public data class BertRuntimeWeights( - val config: BertModelConfig, - // Embeddings - val wordEmbeddings: Tensor, - val positionEmbeddings: Tensor, - val tokenTypeEmbeddings: Tensor, - val embeddingLayerNormWeight: Tensor, - val embeddingLayerNormBias: Tensor, - // Encoder layers - val layers: List>, - // Optional pooler dense (cls token projection) - val poolerDenseWeight: Tensor? = null, - val poolerDenseBias: Tensor? = null, - // Optional sentence-transformers dense projection - val projectionWeight: Tensor? = null, - val projectionBias: Tensor? = null -) diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertWeightLoader.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertWeightLoader.kt deleted file mode 100644 index 6a19b670..00000000 --- a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/BertWeightLoader.kt +++ /dev/null @@ -1,199 +0,0 @@ -package sk.ainet.models.bert - -import sk.ainet.context.ExecutionContext -import sk.ainet.io.ParametersLoader -import sk.ainet.lang.tensor.Shape -import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.types.DType -import kotlin.reflect.KClass - -/** - * Canonical HuggingFace tensor name mappings for BERT. - */ -public object BertTensorNames { - // Embeddings - public const val WORD_EMBEDDINGS: String = "bert.embeddings.word_embeddings.weight" - public const val POSITION_EMBEDDINGS: String = "bert.embeddings.position_embeddings.weight" - public const val TOKEN_TYPE_EMBEDDINGS: String = "bert.embeddings.token_type_embeddings.weight" - public const val EMBEDDING_LN_WEIGHT: String = "bert.embeddings.LayerNorm.weight" - public const val EMBEDDING_LN_BIAS: String = "bert.embeddings.LayerNorm.bias" - - // Per-layer attention - public fun queryWeight(layer: Int): String = "bert.encoder.layer.$layer.attention.self.query.weight" - public fun queryBias(layer: Int): String = "bert.encoder.layer.$layer.attention.self.query.bias" - public fun keyWeight(layer: Int): String = "bert.encoder.layer.$layer.attention.self.key.weight" - public fun keyBias(layer: Int): String = "bert.encoder.layer.$layer.attention.self.key.bias" - public fun valueWeight(layer: Int): String = "bert.encoder.layer.$layer.attention.self.value.weight" - public fun valueBias(layer: Int): String = "bert.encoder.layer.$layer.attention.self.value.bias" - - // Per-layer attention output - public fun attnOutputWeight(layer: Int): String = "bert.encoder.layer.$layer.attention.output.dense.weight" - public fun attnOutputBias(layer: Int): String = "bert.encoder.layer.$layer.attention.output.dense.bias" - public fun attnLayerNormWeight(layer: Int): String = "bert.encoder.layer.$layer.attention.output.LayerNorm.weight" - public fun attnLayerNormBias(layer: Int): String = "bert.encoder.layer.$layer.attention.output.LayerNorm.bias" - - // Per-layer FFN - public fun intermediateWeight(layer: Int): String = "bert.encoder.layer.$layer.intermediate.dense.weight" - public fun intermediateBias(layer: Int): String = "bert.encoder.layer.$layer.intermediate.dense.bias" - public fun outputWeight(layer: Int): String = "bert.encoder.layer.$layer.output.dense.weight" - public fun outputBias(layer: Int): String = "bert.encoder.layer.$layer.output.dense.bias" - public fun outputLayerNormWeight(layer: Int): String = "bert.encoder.layer.$layer.output.LayerNorm.weight" - public fun outputLayerNormBias(layer: Int): String = "bert.encoder.layer.$layer.output.LayerNorm.bias" - - // Pooler - public const val POOLER_DENSE_WEIGHT: String = "bert.pooler.dense.weight" - public const val POOLER_DENSE_BIAS: String = "bert.pooler.dense.bias" - - // Sentence-transformers projection (2_Dense module) - // In the separate 2_Dense/model.safetensors file, keys are "linear.weight"/"linear.bias" - public const val PROJECTION_WEIGHT: String = "linear.weight" - public const val PROJECTION_BIAS: String = "linear.bias" -} - -/** - * Maps a flat tensor name→tensor map (from any source) to typed [BertRuntimeWeights]. - * Format-agnostic: works with SafeTensors, GGUF, ONNX, etc. - */ -public object BertWeightMapper { - - /** - * Prefix used by some HuggingFace checkpoints. Sentence-transformers models - * often store encoder weights without the `bert.` prefix, so we try both. - */ - private const val BERT_PREFIX = "bert." - - public fun map( - tensors: Map>, - config: BertModelConfig - ): BertRuntimeWeights { - val h = config.hiddenSize - val inter = config.intermediateSize - - fun get(name: String): Tensor { - tensors[name]?.let { return it } - // Fallback: try without "bert." prefix (sentence-transformers layout) - if (name.startsWith(BERT_PREFIX)) { - tensors[name.removePrefix(BERT_PREFIX)]?.let { return it } - } - error("Missing BERT tensor: $name (also tried ${name.removePrefix(BERT_PREFIX)})") - } - - fun getOptional(name: String): Tensor? { - tensors[name]?.let { return it } - if (name.startsWith(BERT_PREFIX)) { - tensors[name.removePrefix(BERT_PREFIX)]?.let { return it } - } - return null - } - - fun Tensor<*, *>.require2D(rows: Int, cols: Int, label: String) { - val expected = Shape(rows, cols) - if (shape != expected) error("$label: expected shape $expected but was $shape") - } - - fun Tensor<*, *>.require1D(size: Int, label: String) { - val expected = Shape(size) - if (shape != expected) error("$label: expected shape $expected but was $shape") - } - - // Embeddings - val wordEmb = get(BertTensorNames.WORD_EMBEDDINGS).apply { - require2D(config.vocabSize, h, "word_embeddings") - } - val posEmb = get(BertTensorNames.POSITION_EMBEDDINGS).apply { - require2D(config.maxPositionEmbeddings, h, "position_embeddings") - } - val tokTypeEmb = get(BertTensorNames.TOKEN_TYPE_EMBEDDINGS).apply { - require2D(config.typeVocabSize, h, "token_type_embeddings") - } - val embLnW = get(BertTensorNames.EMBEDDING_LN_WEIGHT).apply { require1D(h, "embedding_ln.weight") } - val embLnB = get(BertTensorNames.EMBEDDING_LN_BIAS).apply { require1D(h, "embedding_ln.bias") } - - // Encoder layers - val layers = (0 until config.numHiddenLayers).map { i -> - BertLayerWeights( - queryWeight = get(BertTensorNames.queryWeight(i)).apply { require2D(h, h, "layer.$i.query.weight") }, - queryBias = get(BertTensorNames.queryBias(i)).apply { require1D(h, "layer.$i.query.bias") }, - keyWeight = get(BertTensorNames.keyWeight(i)).apply { require2D(h, h, "layer.$i.key.weight") }, - keyBias = get(BertTensorNames.keyBias(i)).apply { require1D(h, "layer.$i.key.bias") }, - valueWeight = get(BertTensorNames.valueWeight(i)).apply { require2D(h, h, "layer.$i.value.weight") }, - valueBias = get(BertTensorNames.valueBias(i)).apply { require1D(h, "layer.$i.value.bias") }, - attnOutputWeight = get(BertTensorNames.attnOutputWeight(i)).apply { require2D(h, h, "layer.$i.attn_output.weight") }, - attnOutputBias = get(BertTensorNames.attnOutputBias(i)).apply { require1D(h, "layer.$i.attn_output.bias") }, - attnLayerNormWeight = get(BertTensorNames.attnLayerNormWeight(i)).apply { require1D(h, "layer.$i.attn_ln.weight") }, - attnLayerNormBias = get(BertTensorNames.attnLayerNormBias(i)).apply { require1D(h, "layer.$i.attn_ln.bias") }, - intermediateWeight = get(BertTensorNames.intermediateWeight(i)).apply { require2D(inter, h, "layer.$i.intermediate.weight") }, - intermediateBias = get(BertTensorNames.intermediateBias(i)).apply { require1D(inter, "layer.$i.intermediate.bias") }, - outputWeight = get(BertTensorNames.outputWeight(i)).apply { require2D(h, inter, "layer.$i.output.weight") }, - outputBias = get(BertTensorNames.outputBias(i)).apply { require1D(h, "layer.$i.output.bias") }, - outputLayerNormWeight = get(BertTensorNames.outputLayerNormWeight(i)).apply { require1D(h, "layer.$i.output_ln.weight") }, - outputLayerNormBias = get(BertTensorNames.outputLayerNormBias(i)).apply { require1D(h, "layer.$i.output_ln.bias") } - ) - } - - // Optional pooler - val poolerW = getOptional(BertTensorNames.POOLER_DENSE_WEIGHT)?.apply { require2D(h, h, "pooler.weight") } - val poolerB = getOptional(BertTensorNames.POOLER_DENSE_BIAS)?.apply { require1D(h, "pooler.bias") } - - // Optional sentence-transformers projection - val projDim = config.projectionDim - val projW = getOptional(BertTensorNames.PROJECTION_WEIGHT)?.apply { - if (projDim != null) require2D(projDim, h, "projection.weight") - } - val projB = getOptional(BertTensorNames.PROJECTION_BIAS)?.apply { - if (projDim != null) require1D(projDim, "projection.bias") - } - - return BertRuntimeWeights( - config = config, - wordEmbeddings = wordEmb, - positionEmbeddings = posEmb, - tokenTypeEmbeddings = tokTypeEmb, - embeddingLayerNormWeight = embLnW, - embeddingLayerNormBias = embLnB, - layers = layers, - poolerDenseWeight = poolerW, - poolerDenseBias = poolerB, - projectionWeight = projW, - projectionBias = projB - ) - } -} - -/** - * Load BERT weights from any [ParametersLoader] (SafeTensors, GGUF, ONNX, etc.). - */ -public suspend fun loadBertWeights( - loader: ParametersLoader, - ctx: ExecutionContext, - dtype: KClass, - config: BertModelConfig -): BertRuntimeWeights { - val tensors = mutableMapOf>() - loader.load(ctx, dtype) { name, tensor -> - tensors[name] = tensor - } - return BertWeightMapper.map(tensors, config) -} - -/** - * Load BERT weights from multiple [ParametersLoader]s and merge them. - * - * Sentence-transformers models store the projection layer in a separate - * `2_Dense/model.safetensors` file. This function loads all provided loaders - * and merges their tensors before mapping to [BertRuntimeWeights]. - */ -public suspend fun loadBertWeights( - loaders: List, - ctx: ExecutionContext, - dtype: KClass, - config: BertModelConfig -): BertRuntimeWeights { - val tensors = mutableMapOf>() - for (loader in loaders) { - loader.load(ctx, dtype) { name, tensor -> - tensors[name] = tensor - } - } - return BertWeightMapper.map(tensors, config) -} diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/PooledExecutionContext.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/PooledExecutionContext.kt index 7f253a5f..7338e56d 100644 --- a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/PooledExecutionContext.kt +++ b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/PooledExecutionContext.kt @@ -15,7 +15,7 @@ import sk.ainet.lang.tensor.scratch.SizeClassedScratchPool * val baseCtx = DirectCpuExecutionContext(tensorDataFactory = memSegFactory) * val pooledCtx = PooledExecutionContext(baseCtx) * - * val runtime = BertRuntime(pooledCtx, weights, FP32::class) + * val runtime = createBertEncoderRuntime(config, tensors, pooledCtx, FP32::class) * * // Each forward acquires + releases scratch buffers in a per-call scope. * val v1 = runtime.encode(tokens1) diff --git a/llm-inference/bert/src/jvmMain/kotlin/sk/ainet/models/bert/java/KBertJava.kt b/llm-inference/bert/src/jvmMain/kotlin/sk/ainet/models/bert/java/KBertJava.kt index 4efe648f..9bd425d3 100644 --- a/llm-inference/bert/src/jvmMain/kotlin/sk/ainet/models/bert/java/KBertJava.kt +++ b/llm-inference/bert/src/jvmMain/kotlin/sk/ainet/models/bert/java/KBertJava.kt @@ -3,18 +3,24 @@ package sk.ainet.models.bert.java import kotlinx.coroutines.runBlocking -import sk.ainet.models.bert.* import sk.ainet.context.DirectCpuExecutionContext -import sk.ainet.models.bert.PooledExecutionContext import sk.ainet.io.JvmRandomAccessSource import sk.ainet.io.safetensors.SafeTensorsParametersLoader import sk.ainet.lang.types.FP32 +import sk.ainet.models.bert.BertConfigParser +import sk.ainet.models.bert.BertEncoderRuntime +import sk.ainet.models.bert.BertNetworkLoader +import sk.ainet.models.bert.HuggingFaceTokenizer +import sk.ainet.models.bert.MDBR_LEAF_IR_CONFIG +import sk.ainet.models.bert.PooledExecutionContext +import sk.ainet.models.bert.createBertEncoderRuntime import java.nio.file.Path import kotlin.io.path.exists import kotlin.io.path.readText /** - * Java-friendly facade for loading and running BERT models. + * Java-friendly facade for loading and running BERT embedding models on the + * DSL path ([sk.ainet.models.bert.bertNetwork] + [BertEncoderRuntime]). * * Example usage from Java: * ```java @@ -28,7 +34,8 @@ public object KBertJava { /** * Load a BERT model from a HuggingFace directory containing - * model.safetensors, vocab.txt, and optionally config.json. + * model.safetensors, vocab.txt, and optionally config.json plus the + * sentence-transformers `2_Dense/` projection head. * * @param modelDir Path to the model directory. * @return A KBertSession that implements AutoCloseable. @@ -43,8 +50,21 @@ public object KBertJava { val vocabPath = modelDir.resolve("vocab.txt") require(vocabPath.exists()) { "vocab.txt not found in $modelDir" } - // Detect config from config.json or use defaults - val config = detectConfig(modelDir) + val denseWeights = modelDir.resolve("2_Dense").resolve("model.safetensors") + val denseConfig = modelDir.resolve("2_Dense").resolve("config.json") + + // Detect config from config.json (+ projection head) or use defaults. + // The legacy facade ignored 2_Dense entirely and silently returned + // unprojected embeddings for models that ship a projection. + val configFile = modelDir.resolve("config.json") + val config = if (configFile.exists()) { + BertConfigParser.parse( + configFile.readText(), + if (denseConfig.exists()) denseConfig.readText() else null, + ) + } else { + MDBR_LEAF_IR_CONFIG + } val tokenizer = HuggingFaceTokenizer.fromVocabTxt(vocabPath.readText()) // Pool scratch buffers across encode() calls — embedding workloads @@ -53,43 +73,28 @@ public object KBertJava { // worse than NoopScratchPool. val ctx = PooledExecutionContext(DirectCpuExecutionContext()) - val ingestion = BertIngestion(ctx, FP32::class, config) - val loader = SafeTensorsParametersLoader( - sourceProvider = { JvmRandomAccessSource.open(safetensorsPath.toString()) }, - onProgress = { _, _, _ -> } - ) + val loaders = buildList { + add( + SafeTensorsParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(safetensorsPath.toString()) }, + onProgress = { _, _, _ -> } + ) + ) + if (denseWeights.exists()) { + add( + SafeTensorsParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(denseWeights.toString()) }, + onProgress = { _, _, _ -> } + ) + ) + } + } - val weights = runBlocking { ingestion.load(loader) } - val runtime = BertRuntime(ctx, weights, FP32::class) + val tensors = runBlocking { BertNetworkLoader.loadWeightTensors(loaders, ctx, FP32::class) } + val runtime = createBertEncoderRuntime(config, tensors, ctx) return KBertSession(runtime, tokenizer) } - - private fun detectConfig(modelDir: Path): BertModelConfig { - val configFile = modelDir.resolve("config.json") - if (!configFile.exists()) return MDBR_LEAF_IR_CONFIG - - val json = configFile.readText() - fun extractInt(key: String, default: Int): Int { - val pattern = Regex("\"$key\"\\s*:\\s*(\\d+)") - return pattern.find(json)?.groupValues?.get(1)?.toIntOrNull() ?: default - } - fun extractDouble(key: String, default: Double): Double { - val pattern = Regex("\"$key\"\\s*:\\s*([\\d.eE\\-+]+)") - return pattern.find(json)?.groupValues?.get(1)?.toDoubleOrNull() ?: default - } - - return BertModelConfig( - vocabSize = extractInt("vocab_size", 30522), - hiddenSize = extractInt("hidden_size", 768), - numHiddenLayers = extractInt("num_hidden_layers", 12), - numAttentionHeads = extractInt("num_attention_heads", 12), - intermediateSize = extractInt("intermediate_size", 3072), - maxPositionEmbeddings = extractInt("max_position_embeddings", 512), - typeVocabSize = extractInt("type_vocab_size", 2), - layerNormEps = extractDouble("layer_norm_eps", 1e-12) - ) - } } /** @@ -98,7 +103,7 @@ public object KBertJava { * Implements AutoCloseable for try-with-resources usage. */ public class KBertSession( - private val runtime: BertRuntime, + private val runtime: BertEncoderRuntime, private val tokenizer: HuggingFaceTokenizer ) : AutoCloseable { @@ -110,8 +115,7 @@ public class KBertSession( */ public fun encode(text: String): FloatArray { val tokOutput = tokenizer.encodeWithMetadata(text) - val embedding = runtime.encode(tokOutput.inputIds, tokOutput.attentionMask, tokOutput.tokenTypeIds) - return embedding.data.copyToFloatArray() + return runtime.encode(tokOutput.inputIds, tokOutput.attentionMask) } /** diff --git a/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertEmbeddingsTest.kt b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertEmbeddingsTest.kt new file mode 100644 index 00000000..8e0cb1d1 --- /dev/null +++ b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertEmbeddingsTest.kt @@ -0,0 +1,226 @@ +package sk.ainet.models.bert + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.math.sqrt +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue +import kotlin.test.assertFailsWith + +/** + * Hand-computed verification of the [BertEmbeddings] additive block and the + * name-resolver/mapping pipeline that populates it. + * + * Uses a 0-encoder-layer network so the model output IS the embeddings-block + * output (word + position + token_type row 0, LayerNorm'd). + */ +class BertEmbeddingsTest { + + private val ctx = DirectCpuExecutionContext() + + private val h = 4 + private val vocab = 10 + private val maxPos = 16 + private val typeVocab = 2 + private val eps = 1e-5 + + private fun config(layers: Int = 0) = BertModelConfig( + vocabSize = vocab, + hiddenSize = h, + numHiddenLayers = layers, + numAttentionHeads = 2, + intermediateSize = 8, + maxPositionEmbeddings = maxPos, + typeVocabSize = typeVocab, + layerNormEps = eps, + projectionDim = null, + ) + + /** Row r = [base*(r+1), base*(r+1)*2, base*(r+1)*3, base*(r+1)*4] — varied so LayerNorm keeps signal. */ + private fun table(rows: Int, base: Float): FloatArray = + FloatArray(rows * h) { idx -> + val r = idx / h + val c = idx % h + base * (r + 1) * (c + 1) + } + + private fun tensor2d(rows: Int, data: FloatArray): Tensor = + ctx.fromFloatArray(Shape(rows, h), FP32::class, data) + + private fun tensor1d(data: FloatArray): Tensor = + ctx.fromFloatArray(Shape(h), FP32::class, data) + + /** Bare sentence-transformers names — exercises the bert.-prefix normalization too. */ + private fun embeddingTensors( + word: FloatArray, + pos: FloatArray, + type: FloatArray, + gamma: FloatArray = FloatArray(h) { 1f }, + beta: FloatArray = FloatArray(h) { 0f }, + ): Map> = mapOf( + "embeddings.word_embeddings.weight" to tensor2d(vocab, word), + "embeddings.position_embeddings.weight" to tensor2d(maxPos, pos), + "embeddings.token_type_embeddings.weight" to tensor2d(typeVocab, type), + "embeddings.LayerNorm.weight" to tensor1d(gamma), + "embeddings.LayerNorm.bias" to tensor1d(beta), + ) + + private fun layerNormRow(x: FloatArray, gamma: FloatArray, beta: FloatArray): FloatArray { + val mean = x.average().toFloat() + val variance = x.map { (it - mean) * (it - mean) }.average().toFloat() + val denom = sqrt(variance + eps.toFloat()) + return FloatArray(x.size) { i -> (x[i] - mean) / denom * gamma[i] + beta[i] } + } + + @Test + fun `embeddings block matches hand-computed word + position + type row0 with LayerNorm`() { + val word = table(vocab, 0.2f) + val pos = table(maxPos, 0.1f) + val type = table(typeVocab, 0.05f) + val gamma = floatArrayOf(1f, 2f, 0.5f, 1.5f) + val beta = floatArrayOf(0.1f, -0.1f, 0f, 0.2f) + + val model = BertNetworkLoader.fromTensorMap( + config(layers = 0), + embeddingTensors(word, pos, type, gamma, beta), + ) + + val tokenIds = intArrayOf(3, 0, 7) + val input = ctx.fromFloatArray( + Shape(tokenIds.size), FP32::class, + FloatArray(tokenIds.size) { tokenIds[it].toFloat() }, + ) + val out = model.forward(input, ctx) + assertEquals(listOf(tokenIds.size, h), out.shape.dimensions.toList()) + + tokenIds.forEachIndexed { position, tokenId -> + val summed = FloatArray(h) { c -> + word[tokenId * h + c] + pos[position * h + c] + type[0 * h + c] + } + val expected = layerNormRow(summed, gamma, beta) + for (c in 0 until h) { + val actual = out.data.get(position, c) + assertTrue( + abs(actual - expected[c]) < 1e-4f, + "pos=$position c=$c expected=${expected[c]} actual=$actual", + ) + } + } + } + + @Test + fun `position rows depend on position not token id`() { + // Position rows must point in per-row-distinct directions: every row of + // table() is a scalar multiple of the same (c+1) vector, which LayerNorm + // would normalize to identical rows regardless of position. + val rotatedPos = FloatArray(maxPos * h) { idx -> + val r = idx / h + val c = idx % h + 0.5f * (((r + c) % h) + 1) + } + val model = BertNetworkLoader.fromTensorMap( + config(layers = 0), + embeddingTensors(table(vocab, 0.2f), rotatedPos, table(typeVocab, 0.05f)), + ) + val input = ctx.fromFloatArray(Shape(3), FP32::class, floatArrayOf(5f, 5f, 5f)) + val out = model.forward(input, ctx) + + // Same token at positions 0 vs 2 must differ (position table contributes). + var diff = 0f + for (c in 0 until h) diff += abs(out.data.get(0, c) - out.data.get(2, c)) + assertTrue(diff > 1e-3f, "identical rows for same token at different positions: position embedding lost") + } + + @Test + fun `full 2-layer network maps all parameters from a bare-named sentence-transformers tensor set`() { + val cfg = config(layers = 2) + val tensors = syntheticFullCheckpoint(cfg, bare = true) + // Must not throw (mapped == total); pooler/projection stay unused. + val model = BertNetworkLoader.fromTensorMap(cfg, tensors) + val input = ctx.fromFloatArray(Shape(2), FP32::class, floatArrayOf(1f, 2f)) + val out = model.forward(input, ctx) + assertEquals(listOf(2, h), out.shape.dimensions.toList()) + } + + @Test + fun `full 2-layer network maps all parameters from a bert-prefixed tensor set`() { + val cfg = config(layers = 2) + val model = BertNetworkLoader.fromTensorMap(cfg, syntheticFullCheckpoint(cfg, bare = false)) + val input = ctx.fromFloatArray(Shape(2), FP32::class, floatArrayOf(1f, 2f)) + assertEquals(listOf(2, h), model.forward(input, ctx).shape.dimensions.toList()) + } + + @Test + fun `missing position table fails mapping loudly`() { + val cfg = config(layers = 0) + val tensors = embeddingTensors(table(vocab, 0.2f), table(maxPos, 0.1f), table(typeVocab, 0.05f)) + .filterKeys { !it.contains("position_embeddings") } + assertFailsWith { + BertNetworkLoader.fromTensorMap(cfg, tensors) + } + } + + @Test + fun `sequence longer than maxPositionEmbeddings is rejected`() { + val model = BertNetworkLoader.fromTensorMap( + config(layers = 0), + embeddingTensors(table(vocab, 0.2f), table(maxPos, 0.1f), table(typeVocab, 0.05f)), + ) + val tooLong = ctx.fromFloatArray( + Shape(maxPos + 1), FP32::class, FloatArray(maxPos + 1) { (it % vocab).toFloat() }, + ) + assertFailsWith { model.forward(tooLong, ctx) } + } + + /** Full HF-shaped tensor set for [cfg], bare or bert.-prefixed, incl. unused pooler. */ + private fun syntheticFullCheckpoint( + cfg: BertModelConfig, + bare: Boolean, + ): Map> { + fun identityLike(rows: Int, cols: Int, scale: Float): Tensor { + val data = FloatArray(rows * cols) { idx -> + val r = idx / cols + val c = idx % cols + if (r % cols == c) scale else 0f + } + return ctx.fromFloatArray(Shape(rows, cols), FP32::class, data) + } + + fun vec(size: Int, v: Float) = ctx.fromFloatArray(Shape(size), FP32::class, FloatArray(size) { v }) + + val m = mutableMapOf>() + val p = if (bare) "" else "bert." + m["${p}embeddings.word_embeddings.weight"] = tensor2d(vocab, table(vocab, 0.2f)) + m["${p}embeddings.position_embeddings.weight"] = tensor2d(maxPos, table(maxPos, 0.1f)) + m["${p}embeddings.token_type_embeddings.weight"] = tensor2d(typeVocab, table(typeVocab, 0.05f)) + m["${p}embeddings.LayerNorm.weight"] = vec(h, 1f) + m["${p}embeddings.LayerNorm.bias"] = vec(h, 0f) + for (l in 0 until cfg.numHiddenLayers) { + val lp = "${p}encoder.layer.$l" + m["$lp.attention.self.query.weight"] = identityLike(h, h, 0.5f) + m["$lp.attention.self.query.bias"] = vec(h, 0f) + m["$lp.attention.self.key.weight"] = identityLike(h, h, 0.5f) + m["$lp.attention.self.key.bias"] = vec(h, 0f) + m["$lp.attention.self.value.weight"] = identityLike(h, h, 0.5f) + m["$lp.attention.self.value.bias"] = vec(h, 0f) + m["$lp.attention.output.dense.weight"] = identityLike(h, h, 0.5f) + m["$lp.attention.output.dense.bias"] = vec(h, 0f) + m["$lp.attention.output.LayerNorm.weight"] = vec(h, 1f) + m["$lp.attention.output.LayerNorm.bias"] = vec(h, 0f) + m["$lp.intermediate.dense.weight"] = identityLike(cfg.intermediateSize, h, 0.5f) + m["$lp.intermediate.dense.bias"] = vec(cfg.intermediateSize, 0f) + m["$lp.output.dense.weight"] = identityLike(h, cfg.intermediateSize, 0.5f) + m["$lp.output.dense.bias"] = vec(h, 0f) + m["$lp.output.LayerNorm.weight"] = vec(h, 1f) + m["$lp.output.LayerNorm.bias"] = vec(h, 0f) + } + // Present in real checkpoints, must land in unusedTensors without failing. + m["${p}pooler.dense.weight"] = identityLike(h, h, 1f) + m["${p}pooler.dense.bias"] = vec(h, 0f) + return m + } +} 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 new file mode 100644 index 00000000..2fa39153 --- /dev/null +++ b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertEncoderRuntimeTest.kt @@ -0,0 +1,190 @@ +package sk.ainet.models.bert + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.weights.WeightTensor +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.math.sqrt +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * Synthetic-weight tests for [BertEncoderRuntime] on the DSL path — ports the + * behavioural cases of the legacy eager-runtime test suite: output shapes, + * L2 normalization, mask-weighted pooling, projection, input sensitivity. + */ +class BertEncoderRuntimeTest { + + private val ctx = DirectCpuExecutionContext() + + private val h = 4 + private val inter = 8 + private val vocab = 10 + private val maxPos = 16 + private val typeVocab = 2 + + private fun config(projDim: Int? = null, layers: Int = 2) = BertModelConfig( + vocabSize = vocab, + hiddenSize = h, + numHiddenLayers = layers, + numAttentionHeads = 2, + intermediateSize = inter, + maxPositionEmbeddings = maxPos, + typeVocabSize = typeVocab, + layerNormEps = 1e-5, + projectionDim = projDim, + ) + + private fun varied(rows: Int, cols: Int, scale: Float): Tensor { + val data = FloatArray(rows * cols) { idx -> + val r = idx / cols + val c = idx % cols + scale * (r + 1) * (((r + c) % cols) + 1) + } + return ctx.fromFloatArray(Shape(rows, cols), FP32::class, data) + } + + private fun identityLike(rows: Int, cols: Int, scale: Float): Tensor { + val data = FloatArray(rows * cols) { idx -> + val r = idx / cols + val c = idx % cols + if (r % cols == c) scale else 0f + } + return ctx.fromFloatArray(Shape(rows, cols), FP32::class, data) + } + + private fun vec(size: Int, v: Float): Tensor = + ctx.fromFloatArray(Shape(size), FP32::class, FloatArray(size) { v }) + + /** WeightTensor list shaped like a bare sentence-transformers checkpoint. */ + private fun syntheticTensors(cfg: BertModelConfig, withProjection: Boolean = false): List> { + val out = mutableListOf>() + fun add(name: String, t: Tensor) { + out += WeightTensor(name, t.shape.dimensions.toList(), t) + } + add("embeddings.word_embeddings.weight", varied(vocab, h, 0.2f)) + add("embeddings.position_embeddings.weight", varied(maxPos, h, 0.1f)) + add("embeddings.token_type_embeddings.weight", varied(typeVocab, h, 0.05f)) + add("embeddings.LayerNorm.weight", vec(h, 1f)) + add("embeddings.LayerNorm.bias", vec(h, 0f)) + for (l in 0 until cfg.numHiddenLayers) { + val p = "encoder.layer.$l" + add("$p.attention.self.query.weight", identityLike(h, h, 0.5f)) + add("$p.attention.self.query.bias", vec(h, 0f)) + add("$p.attention.self.key.weight", identityLike(h, h, 0.5f)) + add("$p.attention.self.key.bias", vec(h, 0f)) + add("$p.attention.self.value.weight", identityLike(h, h, 0.5f)) + add("$p.attention.self.value.bias", vec(h, 0f)) + add("$p.attention.output.dense.weight", identityLike(h, h, 0.5f)) + add("$p.attention.output.dense.bias", vec(h, 0f)) + add("$p.attention.output.LayerNorm.weight", vec(h, 1f)) + add("$p.attention.output.LayerNorm.bias", vec(h, 0f)) + add("$p.intermediate.dense.weight", identityLike(inter, h, 0.5f)) + add("$p.intermediate.dense.bias", vec(inter, 0f)) + add("$p.output.dense.weight", identityLike(h, inter, 0.5f)) + add("$p.output.dense.bias", vec(h, 0f)) + add("$p.output.LayerNorm.weight", vec(h, 1f)) + add("$p.output.LayerNorm.bias", vec(h, 0f)) + } + if (withProjection) { + add(BertNetworkLoader.PROJECTION_WEIGHT, identityLike(cfg.projectionDim!!, h, 1f)) + add(BertNetworkLoader.PROJECTION_BIAS, vec(cfg.projectionDim!!, 0f)) + } + return BertNetworkLoader.normalizeTensorNames(out) + } + + private fun runtime(cfg: BertModelConfig = config(), withProjection: Boolean = false): BertEncoderRuntime = + createBertEncoderRuntime(cfg, syntheticTensors(cfg, withProjection), ctx) + + @Test + fun forward_producesCorrectShape() { + val out = runtime().forward(intArrayOf(1, 2, 3)) + assertEquals(2, out.rank) + assertEquals(3, out.shape[0]) + assertEquals(h, out.shape[1]) + } + + @Test + fun forward_singleToken() { + assertEquals(Shape(1, h), runtime().forward(intArrayOf(0)).shape) + } + + @Test + fun encode_producesVectorOfRuntimeDimensions() { + val rt = runtime() + assertEquals(h, rt.dimensions) + assertEquals(h, rt.encode(intArrayOf(1, 2, 3)).size) + } + + @Test + fun encode_isL2Normalized() { + val e = runtime().encode(intArrayOf(1, 2, 3)) + val norm = sqrt(e.sumOf { (it * it).toDouble() }).toFloat() + assertTrue(abs(norm - 1f) < 0.01f, "L2 norm should be ~1.0, got $norm") + } + + @Test + fun encode_withAttentionMask_normalizedAndDiffersFromUnmasked() { + val rt = runtime() + val masked = rt.encode(intArrayOf(1, 2, 0), attentionMask = intArrayOf(1, 1, 0)) + val unmasked = rt.encode(intArrayOf(1, 2, 0)) + + val norm = sqrt(masked.sumOf { (it * it).toDouble() }).toFloat() + assertTrue(abs(norm - 1f) < 0.01f, "L2 norm should be ~1.0 with mask, got $norm") + + val diff = masked.indices.sumOf { abs(masked[it] - unmasked[it]).toDouble() } + assertTrue(diff > 1e-4, "mask must change pooling: diff=$diff") + } + + @Test + fun encode_mismatchedMaskIsRejected() { + assertFailsWith { + runtime().encode(intArrayOf(1, 2, 3), attentionMask = intArrayOf(1, 1)) + } + } + + @Test + fun forward_differentInputsProduceDifferentOutputs() { + val rt = runtime() + val a = rt.forward(intArrayOf(1, 2, 3)) + val b = rt.forward(intArrayOf(4, 5, 6)) + var allSame = true + outer@ for (i in 0 until a.shape[0]) { + for (j in 0 until a.shape[1]) { + if (a.data.get(i, j) != b.data.get(i, j)) { + allSame = false + break@outer + } + } + } + assertTrue(!allSame, "Different inputs should produce different outputs") + } + + @Test + fun encode_withProjection_outputsProjectionDim() { + val projDim = 2 + val rt = runtime(config(projDim = projDim, layers = 1), withProjection = true) + assertEquals(projDim, rt.dimensions) + val e = rt.encode(intArrayOf(1, 2, 3)) + assertEquals(projDim, e.size) + val norm = sqrt(e.sumOf { (it * it).toDouble() }).toFloat() + assertTrue(abs(norm - 1f) < 0.01f, "projected embedding must stay L2-normalized, got $norm") + } + + @Test + fun createRuntime_missingProjectionTensorsFailsWhenConfigDeclaresProjection() { + val cfg = config(projDim = 2, layers = 1) + assertFailsWith { + createBertEncoderRuntime(cfg, syntheticTensors(cfg, withProjection = false), ctx) + } + } + + @Test + fun encode_emptyInputIsRejected() { + assertFailsWith { runtime().encode(intArrayOf()) } + } +} diff --git a/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertNumericalAccuracyTest.kt b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertNumericalAccuracyTest.kt index d1455b19..9c5fcf98 100644 --- a/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertNumericalAccuracyTest.kt +++ b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertNumericalAccuracyTest.kt @@ -266,8 +266,8 @@ class BertNumericalAccuracyTest { val loader = SafeTensorsParametersLoader( sourceProvider = { JvmRandomAccessSource.open(mainSafetensors.toString()) } ) - val weights = loadBertWeights(loader, ctx, FP32::class, encoderConfig) - val runtime = BertRuntime(ctx, weights, FP32::class) + val tensors = BertNetworkLoader.loadWeightTensors(listOf(loader), ctx, FP32::class) + val runtime = createBertEncoderRuntime(encoderConfig, tensors, ctx) for (input in reference.inputs) { val tokenIds = input.inputIds.toIntArray() @@ -319,18 +319,19 @@ class BertNumericalAccuracyTest { sourceProvider = { JvmRandomAccessSource.open(denseSafetensors.toString()) } ) ) - val weights = loadBertWeights(loaders, ctx, FP32::class, config) - val runtime = BertRuntime(ctx, weights, FP32::class) + val tensors = BertNetworkLoader.loadWeightTensors(loaders, ctx, FP32::class) + val runtime = createBertEncoderRuntime(config, tensors, ctx) val embeddings = mutableListOf>() for (input in reference.inputs) { val tokenIds = input.inputIds.toIntArray() val attentionMask = input.attentionMask.toIntArray() - val tokenTypeIds = input.tokenTypeIds.toIntArray() - val embedding = runtime.encode(tokenIds, attentionMask, tokenTypeIds) - val actual = embedding.toFloatList() + // Reference inputs are single-segment (token_type_ids all zero), + // which is exactly what the DSL network encodes. + check(input.tokenTypeIds.all { it == 0 }) { "reference input is not single-segment" } + val actual = runtime.encode(tokenIds, attentionMask).toList() embeddings.add(actual) assertEquals( diff --git a/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertRuntimeTest.kt b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertRuntimeTest.kt deleted file mode 100644 index 8896ed4e..00000000 --- a/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertRuntimeTest.kt +++ /dev/null @@ -1,237 +0,0 @@ -package sk.ainet.models.bert - -import sk.ainet.context.DirectCpuExecutionContext -import sk.ainet.lang.tensor.Shape -import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.types.FP32 -import kotlin.test.Test -import kotlin.test.assertEquals -import kotlin.test.assertTrue - -class BertRuntimeTest { - - private val ctx = DirectCpuExecutionContext() - - /** Embedding matrix with variation across the hidden dimension so LayerNorm doesn't zero it out. */ - private fun variedEmbedding(rows: Int, cols: Int, scale: Float = 0.1f): Tensor { - val data = FloatArray(rows * cols) { idx -> - val row = idx / cols - val col = idx % cols - scale * (row + 1) * (col + 1).toFloat() / cols - } - return ctx.fromFloatArray(Shape(rows, cols), FP32::class, data) - } - - /** Identity-like matrix (scaled) that preserves per-dimension variation through linear projections. */ - private fun identityLike(rows: Int, cols: Int, scale: Float = 1f): Tensor { - val data = FloatArray(rows * cols) { idx -> - val r = idx / cols - val c = idx % cols - if (r % cols == c) scale else 0f - } - return ctx.fromFloatArray(Shape(rows, cols), FP32::class, data) - } - - private fun ones1D(size: Int) = ctx.full(Shape(size), FP32::class, 1f) - private fun zeros1D(size: Int) = ctx.full(Shape(size), FP32::class, 0f) - - private fun buildLayerWeights(h: Int, inter: Int): BertLayerWeights = - BertLayerWeights( - queryWeight = identityLike(h, h, 0.5f), - queryBias = zeros1D(h), - keyWeight = identityLike(h, h, 0.5f), - keyBias = zeros1D(h), - valueWeight = identityLike(h, h, 0.5f), - valueBias = zeros1D(h), - attnOutputWeight = identityLike(h, h, 0.5f), - attnOutputBias = zeros1D(h), - attnLayerNormWeight = ones1D(h), - attnLayerNormBias = zeros1D(h), - intermediateWeight = identityLike(inter, h, 0.5f), - intermediateBias = zeros1D(inter), - outputWeight = identityLike(h, inter, 0.5f), - outputBias = zeros1D(h), - outputLayerNormWeight = ones1D(h), - outputLayerNormBias = zeros1D(h) - ) - - /** - * Build a tiny BERT model (2 layers, 4 hidden dim, 2 heads, 8 intermediate) - * with varied embeddings and identity-like projections for smoke testing. - */ - private fun buildTinyModel(): Pair, BertModelConfig> { - val h = 4 - val inter = 8 - val numHeads = 2 - val numLayers = 2 - val vocabSize = 10 - val maxPos = 16 - val typeVocab = 2 - - val config = BertModelConfig( - vocabSize = vocabSize, - hiddenSize = h, - numHiddenLayers = numLayers, - numAttentionHeads = numHeads, - intermediateSize = inter, - maxPositionEmbeddings = maxPos, - typeVocabSize = typeVocab, - layerNormEps = 1e-5, - projectionDim = null - ) - - val layers = (0 until numLayers).map { buildLayerWeights(h, inter) } - - val weights = BertRuntimeWeights( - config = config, - wordEmbeddings = variedEmbedding(vocabSize, h, 0.2f), - positionEmbeddings = variedEmbedding(maxPos, h, 0.1f), - tokenTypeEmbeddings = variedEmbedding(typeVocab, h, 0.05f), - embeddingLayerNormWeight = ones1D(h), - embeddingLayerNormBias = zeros1D(h), - layers = layers - ) - - val runtime = BertRuntime(ctx, weights, FP32::class) - return runtime to config - } - - @Test - fun forward_producesCorrectShape() { - val (runtime, config) = buildTinyModel() - val tokenIds = intArrayOf(1, 2, 3) // 3 tokens - val output = runtime.forward(tokenIds) - - assertEquals(2, output.rank, "Output should be 2D") - assertEquals(3, output.shape[0], "Sequence length should be 3") - assertEquals(config.hiddenSize, output.shape[1], "Hidden dim should match config") - } - - @Test - fun forward_singleToken() { - val (runtime, config) = buildTinyModel() - val output = runtime.forward(intArrayOf(0)) - - assertEquals(Shape(1, config.hiddenSize), output.shape) - } - - @Test - fun encode_producesVector() { - val (runtime, config) = buildTinyModel() - val embedding = runtime.encode(intArrayOf(1, 2, 3)) - - // Should be a 1D vector of size hiddenSize (no projection configured) - assertEquals(1, embedding.rank, "Embedding should be 1D") - assertEquals(config.hiddenSize, embedding.shape[0], "Embedding dim should match hiddenSize") - } - - @Test - fun encode_isL2Normalized() { - val (runtime, _) = buildTinyModel() - val embedding = runtime.encode(intArrayOf(1, 2, 3)) - - // L2 norm should be ~1.0 - var sumSq = 0f - for (i in 0 until embedding.shape[0]) { - val v = embedding.data[i] as Float - sumSq += v * v - } - val norm = kotlin.math.sqrt(sumSq) - assertTrue(kotlin.math.abs(norm - 1f) < 0.01f, "L2 norm should be ~1.0, got $norm") - } - - @Test - fun encode_withAttentionMask() { - val (runtime, _) = buildTinyModel() - // Encode with 3 tokens, but mask says only first 2 are real - val embedding = runtime.encode( - tokenIds = intArrayOf(1, 2, 0), - attentionMask = intArrayOf(1, 1, 0) - ) - - assertEquals(1, embedding.rank, "Embedding should be 1D") - // Should still produce valid normalized output - var sumSq = 0f - for (i in 0 until embedding.shape[0]) { - val v = embedding.data[i] as Float - sumSq += v * v - } - val norm = kotlin.math.sqrt(sumSq) - assertTrue(kotlin.math.abs(norm - 1f) < 0.01f, "L2 norm should be ~1.0 with mask, got $norm") - } - - @Test - fun encode_withTokenTypeIds() { - val (runtime, _) = buildTinyModel() - val embedding = runtime.encode( - tokenIds = intArrayOf(1, 2, 3, 4), - tokenTypeIds = intArrayOf(0, 0, 1, 1) - ) - - assertEquals(1, embedding.rank) - } - - @Test - fun forward_differentInputsProduceDifferentOutputs() { - val (runtime, _) = buildTinyModel() - val out1 = runtime.forward(intArrayOf(1, 2, 3)) - val out2 = runtime.forward(intArrayOf(4, 5, 6)) - - // Different inputs should generally produce different hidden states - var allSame = true - for (i in 0 until out1.shape[0]) { - for (j in 0 until out1.shape[1]) { - if (out1.data[i, j] != out2.data[i, j]) { - allSame = false - break - } - } - } - assertTrue(!allSame, "Different inputs should produce different outputs") - } - - @Test - fun encode_withProjection() { - val h = 4 - val projDim = 2 - val config = BertModelConfig( - vocabSize = 10, - hiddenSize = h, - numHiddenLayers = 1, - numAttentionHeads = 2, - intermediateSize = 8, - maxPositionEmbeddings = 16, - typeVocabSize = 2, - layerNormEps = 1e-5, - projectionDim = projDim - ) - - val weights = BertRuntimeWeights( - config = config, - wordEmbeddings = variedEmbedding(10, h, 0.2f), - positionEmbeddings = variedEmbedding(16, h, 0.1f), - tokenTypeEmbeddings = variedEmbedding(2, h, 0.05f), - embeddingLayerNormWeight = ones1D(h), - embeddingLayerNormBias = zeros1D(h), - layers = listOf(buildLayerWeights(h, 8)), - projectionWeight = identityLike(projDim, h, 0.3f), - projectionBias = zeros1D(projDim) - ) - - val runtime = BertRuntime(ctx, weights, FP32::class) - val embedding = runtime.encode(intArrayOf(1, 2, 3)) - - // Projection should reduce to projDim - assertEquals(1, embedding.rank) - assertEquals(projDim, embedding.shape[0]) - - // Still L2 normalized - var sumSq = 0f - for (i in 0 until embedding.shape[0]) { - val v = embedding.data[i] as Float - sumSq += v * v - } - val norm = kotlin.math.sqrt(sumSq) - assertTrue(kotlin.math.abs(norm - 1f) < 0.01f, "L2 norm should be ~1.0 with projection, got $norm") - } -} diff --git a/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertStableHloExportTest.kt b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertStableHloExportTest.kt new file mode 100644 index 00000000..91c32655 --- /dev/null +++ b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertStableHloExportTest.kt @@ -0,0 +1,41 @@ +package sk.ainet.models.bert + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.compile.hlo.toStableHlo +import java.io.File +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * StableHLO export smoke test: the DSL BERT encoder traces to a ComputeGraph + * that lowers to StableHLO text (the entry point of the MLIR/IREE path). + * Compilation/execution with iree-compile is out of scope here — this gates + * that the network definition stays export-clean. + */ +class BertStableHloExportTest { + + private val ctx = DirectCpuExecutionContext() + + @Test + fun bertEncoderExportsToStableHlo() { + val cfg = BertSyntheticFixtures.tinyConfig(layers = 1) + val tensors = BertSyntheticFixtures.weightTensors(ctx, cfg) + val runtime = createBertEncoderRuntime(cfg, tensors, ctx) + + val tape = runtime.exportTape(seqLen = 4) + val graph = tape.toComputeGraph(synthesizeExternalInputs = true) + val mlir = toStableHlo(graph, "bert_encoder").content + + assertTrue(mlir.contains("module"), "no MLIR module emitted") + assertTrue(mlir.contains("func.func"), "no function emitted") + // The ops that carry BERT's structure must survive lowering: + assertTrue(mlir.contains("stablehlo.gather") || mlir.contains("\"stablehlo.gather\""), "no gather (word embeddings) in export") + assertTrue(mlir.contains("stablehlo.dot") || mlir.contains("stablehlo.dot_general"), "no matmul in export") + assertTrue(mlir.contains("stablehlo.add"), "no add in export") + + val out = File(System.getProperty("bertMlirOut") ?: "build/build-mlir/bert_encoder.mlir") + out.parentFile?.mkdirs() + out.writeText(mlir) + println("WROTE_MLIR ${out.absolutePath} (${mlir.lines().size} lines)") + } +} diff --git a/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertSyntheticFixtures.kt b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertSyntheticFixtures.kt new file mode 100644 index 00000000..a2fc8f11 --- /dev/null +++ b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertSyntheticFixtures.kt @@ -0,0 +1,97 @@ +package sk.ainet.models.bert + +import sk.ainet.context.ExecutionContext +import sk.ainet.io.weights.WeightTensor +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 + +/** + * Deterministic synthetic BERT checkpoints for tests: varied (non-degenerate + * under LayerNorm) embeddings and non-trivial but well-conditioned projections. + */ +internal object BertSyntheticFixtures { + + fun tinyConfig( + h: Int = 4, + layers: Int = 2, + projDim: Int? = null, + ) = BertModelConfig( + vocabSize = 10, + hiddenSize = h, + numHiddenLayers = layers, + numAttentionHeads = 2, + intermediateSize = h * 2, + maxPositionEmbeddings = 16, + typeVocabSize = 2, + layerNormEps = 1e-5, + projectionDim = projDim, + ) + + fun weightTensors( + ctx: ExecutionContext, + cfg: BertModelConfig, + withProjection: Boolean = false, + ): List> { + val h = cfg.hiddenSize + + fun varied(rows: Int, cols: Int, scale: Float): Tensor { + val data = FloatArray(rows * cols) { idx -> + val r = idx / cols + val c = idx % cols + scale * (r + 1) * (((r + c) % cols) + 1) + } + return ctx.fromFloatArray(Shape(rows, cols), FP32::class, data) + } + + fun mixed(rows: Int, cols: Int, scale: Float): Tensor { + // Diagonal-dominant with small off-diagonal texture: keeps values + // bounded through 2 layers but exercises real mixing. + val data = FloatArray(rows * cols) { idx -> + val r = idx / cols + val c = idx % cols + val diag = if (r % cols == c) scale else 0f + diag + 0.05f * (((r * 3 + c * 7) % 5) - 2) + } + return ctx.fromFloatArray(Shape(rows, cols), FP32::class, data) + } + + fun vec(size: Int, v: Float): Tensor = + ctx.fromFloatArray(Shape(size), FP32::class, FloatArray(size) { v }) + + val out = mutableListOf>() + fun add(name: String, t: Tensor) { + out += WeightTensor(name, t.shape.dimensions.toList(), t) + } + + add("embeddings.word_embeddings.weight", varied(cfg.vocabSize, h, 0.2f)) + add("embeddings.position_embeddings.weight", varied(cfg.maxPositionEmbeddings, h, 0.1f)) + add("embeddings.token_type_embeddings.weight", varied(cfg.typeVocabSize, h, 0.05f)) + add("embeddings.LayerNorm.weight", vec(h, 1f)) + add("embeddings.LayerNorm.bias", vec(h, 0f)) + for (l in 0 until cfg.numHiddenLayers) { + val p = "encoder.layer.$l" + add("$p.attention.self.query.weight", mixed(h, h, 0.5f)) + add("$p.attention.self.query.bias", vec(h, 0.01f)) + add("$p.attention.self.key.weight", mixed(h, h, 0.5f)) + add("$p.attention.self.key.bias", vec(h, 0f)) + add("$p.attention.self.value.weight", mixed(h, h, 0.5f)) + add("$p.attention.self.value.bias", vec(h, 0f)) + add("$p.attention.output.dense.weight", mixed(h, h, 0.5f)) + add("$p.attention.output.dense.bias", vec(h, 0f)) + add("$p.attention.output.LayerNorm.weight", vec(h, 1f)) + add("$p.attention.output.LayerNorm.bias", vec(h, 0f)) + add("$p.intermediate.dense.weight", mixed(cfg.intermediateSize, h, 0.5f)) + add("$p.intermediate.dense.bias", vec(cfg.intermediateSize, 0f)) + add("$p.output.dense.weight", mixed(h, cfg.intermediateSize, 0.5f)) + add("$p.output.dense.bias", vec(h, 0f)) + add("$p.output.LayerNorm.weight", vec(h, 1f)) + add("$p.output.LayerNorm.bias", vec(h, 0f)) + } + if (withProjection) { + val projDim = requireNotNull(cfg.projectionDim) + add(BertNetworkLoader.PROJECTION_WEIGHT, mixed(projDim, h, 1f)) + } + return BertNetworkLoader.normalizeTensorNames(out) + } +} diff --git a/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertTraceParityTest.kt b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertTraceParityTest.kt new file mode 100644 index 00000000..b6013579 --- /dev/null +++ b/llm-inference/bert/src/jvmTest/kotlin/sk/ainet/models/bert/BertTraceParityTest.kt @@ -0,0 +1,114 @@ +package sk.ainet.models.bert + +import kotlinx.coroutines.test.runTest +import org.junit.jupiter.api.Assumptions.assumeTrue +import org.junit.jupiter.api.Tag +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.safetensors.SafeTensorsParametersLoader +import sk.ainet.lang.types.FP32 +import java.nio.file.Path +import kotlin.io.path.exists +import kotlin.io.path.isDirectory +import kotlin.io.path.readText +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * DIRECT vs OPTIMIZED parity: the traced/fused ComputeGraph must reproduce + * the eager module execution — on a synthetic tiny model (always) and on the + * real LEAF checkpoint (integration, self-skipping). + */ +class BertTraceParityTest { + + private val ctx = DirectCpuExecutionContext() + + private fun maxDiff(a: FloatArray, b: FloatArray): Double { + assertEquals(a.size, b.size) + var d = 0.0 + for (i in a.indices) { + val diff = abs((a[i] - b[i]).toDouble()) + if (diff > d) d = diff + } + return d + } + + @Test + fun syntheticModel_optimizedMatchesDirect() { + val cfg = BertSyntheticFixtures.tinyConfig(layers = 2) + val tensors = BertSyntheticFixtures.weightTensors(ctx, cfg) + val direct = createBertEncoderRuntime(cfg, tensors, ctx, mode = BertExecutionMode.DIRECT) + val optimized = createBertEncoderRuntime(cfg, tensors, ctx, mode = BertExecutionMode.OPTIMIZED) + + val sequences = listOf( + intArrayOf(1, 2, 3), + intArrayOf(4, 5, 6, 7, 8), + intArrayOf(9, 0, 1), // same length as the first — exercises the compiled cache hit + ) + for (tokens in sequences) { + val d = maxDiff(direct.encode(tokens), optimized.encode(tokens)) + assertTrue(d < 1e-4, "DIRECT vs OPTIMIZED encode diff $d for ${tokens.size} tokens") + } + } + + @Test + fun syntheticModel_compileReportsSingleInput() { + val cfg = BertSyntheticFixtures.tinyConfig(layers = 1) + val tensors = BertSyntheticFixtures.weightTensors(ctx, cfg) + val optimized = createBertEncoderRuntime(cfg, tensors, ctx, mode = BertExecutionMode.OPTIMIZED) + val diagnostics = optimized.compile(sampleSeqLen = 4) + assertTrue(diagnostics.any { it.startsWith("Traced graph:") }, "missing trace diagnostics: $diagnostics") + assertTrue(diagnostics.any { it.startsWith("Input node:") }, "missing input-node diagnostics: $diagnostics") + } + + @Test + @Tag("integration") + fun realModel_optimizedMatchesDirect() = runTest { + val modelDir = resolveModelDir() + assumeTrue(modelDir != null, "No LEAF model available — skipping") + modelDir!! + val mainFile = modelDir.resolve("model.safetensors") + assumeTrue(mainFile.exists(), "model.safetensors not found") + + val denseFile = modelDir.resolve("2_Dense").resolve("model.safetensors") + val denseConfig = modelDir.resolve("2_Dense").resolve("config.json") + val config = BertConfigParser.parse( + modelDir.resolve("config.json").readText(), + if (denseConfig.exists()) denseConfig.readText() else null, + ) + + val loaders = buildList { + add(SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(mainFile.toString()) })) + if (denseFile.exists()) { + add(SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(denseFile.toString()) })) + } + } + val tensors = BertNetworkLoader.loadWeightTensors(loaders, ctx, FP32::class) + val direct = createBertEncoderRuntime(config, tensors, ctx, mode = BertExecutionMode.DIRECT) + val optimized = createBertEncoderRuntime(config, tensors, ctx, mode = BertExecutionMode.OPTIMIZED) + + val sequences = listOf( + intArrayOf(101, 7592, 2088, 102), + intArrayOf(101, 1996, 4248, 2829, 4419, 14523, 2058, 1996, 13971, 3899, 102), + ) + for (tokens in sequences) { + val d = maxDiff(direct.encode(tokens), optimized.encode(tokens)) + assertTrue(d < 1e-4, "real-model DIRECT vs OPTIMIZED encode diff $d for ${tokens.size} tokens") + println("trace parity ok (${tokens.size} tokens): max diff $d") + } + } + + private fun resolveModelDir(): Path? { + System.getenv("LEAF_MODEL_DIR")?.let { p -> + Path.of(p).takeIf { it.isDirectory() }?.let { return it } + } + val home = System.getProperty("user.home") ?: return null + Path.of(home, ".cache", "skainet", "models", "MongoDB_mdbr-leaf-mt") + .takeIf { it.isDirectory() }?.let { return it } + val snapshots = Path.of(home, ".cache", "huggingface", "hub", "models--MongoDB--mdbr-leaf-ir", "snapshots") + if (!snapshots.isDirectory()) return null + return snapshots.toFile().listFiles()?.firstOrNull { it.isDirectory }?.toPath() + } +} diff --git a/llm-providers/api/llm-providers.api b/llm-providers/api/llm-providers.api index 36861f7c..7d502d39 100644 --- a/llm-providers/api/llm-providers.api +++ b/llm-providers/api/llm-providers.api @@ -1,3 +1,17 @@ +public final class sk/ainet/llm/providers/BertEmbeddingModel { + public static final field INSTANCE Lsk/ainet/llm/providers/BertEmbeddingModel; + public static final fun fromHuggingFace (Ljava/lang/String;)Lsk/ainet/llm/api/EmbeddingModel; + public static final fun fromHuggingFace (Ljava/lang/String;Ljava/lang/String;)Lsk/ainet/llm/api/EmbeddingModel; + public static final fun fromHuggingFace (Ljava/lang/String;Ljava/lang/String;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/llm/api/EmbeddingModel; + public static final fun fromHuggingFace (Ljava/lang/String;Ljava/lang/String;Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/bert/BertExecutionMode;)Lsk/ainet/llm/api/EmbeddingModel; + public static synthetic fun fromHuggingFace$default (Ljava/lang/String;Ljava/lang/String;Lsk/ainet/context/ExecutionContext;Lsk/ainet/models/bert/BertExecutionMode;ILjava/lang/Object;)Lsk/ainet/llm/api/EmbeddingModel; + public static final fun fromSafeTensors (Ljava/nio/file/Path;)Lsk/ainet/llm/api/EmbeddingModel; + 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 final class sk/ainet/llm/providers/SkaiNetChatModel : sk/ainet/llm/api/StreamingChatModel { public fun (Lsk/ainet/apps/llm/InferenceRuntime;Lsk/ainet/apps/llm/Tokenizer;Lsk/ainet/apps/kllama/chat/ChatTemplate;Lsk/ainet/llm/api/ChatOptions;ILjava/util/Set;Lkotlin/random/Random;Ljava/lang/String;)V public synthetic fun (Lsk/ainet/apps/llm/InferenceRuntime;Lsk/ainet/apps/llm/Tokenizer;Lsk/ainet/apps/kllama/chat/ChatTemplate;Lsk/ainet/llm/api/ChatOptions;ILjava/util/Set;Lkotlin/random/Random;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -8,8 +22,8 @@ public final class sk/ainet/llm/providers/SkaiNetChatModel : sk/ainet/llm/api/St } public final class sk/ainet/llm/providers/SkaiNetEmbeddingModel : sk/ainet/llm/api/EmbeddingModel { - public fun (Lsk/ainet/models/bert/BertRuntime;Lsk/ainet/apps/llm/Tokenizer;ILjava/lang/String;)V - public synthetic fun (Lsk/ainet/models/bert/BertRuntime;Lsk/ainet/apps/llm/Tokenizer;ILjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lsk/ainet/models/bert/BertEncoderRuntime;Lsk/ainet/apps/llm/Tokenizer;ILjava/lang/String;)V + public synthetic fun (Lsk/ainet/models/bert/BertEncoderRuntime;Lsk/ainet/apps/llm/Tokenizer;ILjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun call (Lsk/ainet/llm/api/EmbeddingRequest;)Lsk/ainet/llm/api/EmbeddingResponse; public fun getDimensions ()I } diff --git a/llm-providers/build.gradle.kts b/llm-providers/build.gradle.kts index b5713d96..7934bdb8 100644 --- a/llm-providers/build.gradle.kts +++ b/llm-providers/build.gradle.kts @@ -29,6 +29,15 @@ dependencies { implementation(libs.kotlinx.coroutines) implementation(libs.kotlinx.serialization.json) + // BertEmbeddingModel one-call factory: local SafeTensors loading + built-in + // Hugging Face download (hf:// URIs). kotlinx-io appears in data-source's + // public API but is declared implementation upstream, so add it explicitly. + implementation(libs.skainet.backend.cpu) + implementation(libs.skainet.io.core) + implementation(libs.skainet.io.safetensors) + implementation(libs.skainet.data.source) + implementation(libs.kotlinx.io.core) + testImplementation(libs.kotlin.test) testImplementation(libs.kotlinx.coroutines.test) testImplementation(libs.junit.jupiter) 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 new file mode 100644 index 00000000..a7c5dadc --- /dev/null +++ b/llm-providers/src/main/kotlin/sk/ainet/llm/providers/BertEmbeddingModel.kt @@ -0,0 +1,217 @@ +package sk.ainet.llm.providers + +import kotlinx.coroutines.runBlocking +import kotlinx.io.buffered +import kotlinx.io.files.SystemFileSystem +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.data.source.CachePolicy +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 +import sk.ainet.models.bert.BertNetworkLoader +import sk.ainet.models.bert.HuggingFaceTokenizer +import sk.ainet.models.bert.PooledExecutionContext +import sk.ainet.models.bert.createBertEncoderRuntime +import java.nio.file.Files +import java.nio.file.Path +import java.nio.file.StandardCopyOption +import kotlin.io.path.createDirectories +import kotlin.io.path.exists +import kotlin.io.path.name +import kotlin.io.path.readText +import kotlinx.io.files.Path as KotlinxPath + +/** + * One-call factories for BERT sentence-embedding models behind the neutral + * [EmbeddingModel] SPI. + * + * ```kotlin + * // From the Hugging Face Hub — downloads and caches on first use: + * val model = BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-mt") + * + * // From a local snapshot directory: + * val model = BertEmbeddingModel.fromSafeTensors(Path.of("/models/mdbr-leaf-mt")) + * + * val vector = model.embed("How do I reset a password?") + * ``` + * + * Auto-detects inside the snapshot: the base weights file + * (`model.safetensors` / `pytorch_model.safetensors` / any `*.safetensors`), + * the optional sentence-transformers projection head + * (`2_Dense/model.safetensors` + `2_Dense/config.json`), the model config + * (`config.json`), and the tokenizer (`vocab.txt`, falling back to + * `tokenizer.json`). + */ +public object BertEmbeddingModel { + + /** + * Load a HuggingFace sentence-transformers BERT snapshot directory. + * + * @param modelDir directory containing the snapshot + * @param ctx execution context; defaults to a pooled CPU context (scratch + * buffers reused across `embed` calls) + * @param mode eager [BertExecutionMode.DIRECT] (default) or traced/fused + * [BertExecutionMode.OPTIMIZED] + * @param modelId reported via [EmbeddingModel] response metadata + */ + @JvmStatic + @JvmOverloads + public fun fromSafeTensors( + modelDir: Path, + ctx: ExecutionContext = PooledExecutionContext(DirectCpuExecutionContext()), + mode: BertExecutionMode = BertExecutionMode.DIRECT, + modelId: String? = modelDir.fileName?.toString(), + ): EmbeddingModel { + val weightsFile = resolveWeightsFile(modelDir) + val denseWeights = modelDir.resolve("2_Dense").resolve("model.safetensors") + val denseConfig = modelDir.resolve("2_Dense").resolve("config.json") + + val configFile = modelDir.resolve("config.json") + require(configFile.exists()) { "config.json not found in $modelDir" } + val config = BertConfigParser.parse( + configFile.readText(), + if (denseConfig.exists()) denseConfig.readText() else null, + ) + + val tokenizer = loadTokenizer(modelDir) + + val loaders = buildList { + add(SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(weightsFile.toString()) })) + if (denseWeights.exists()) { + add(SafeTensorsParametersLoader(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) + + return SkaiNetEmbeddingModel( + runtime = runtime, + tokenizer = tokenizer, + modelId = modelId, + ) + } + + /** + * Load a BERT sentence-transformers model straight from the Hugging Face + * Hub, downloading missing snapshot files into + * `~/.cache/skainet/models/_` on first use (offline-safe + * afterwards). Gated repos pick up `HF_TOKEN` / `HUGGING_FACE_HUB_TOKEN`. + * + * @param repoId Hub repo, e.g. `"MongoDB/mdbr-leaf-mt"` + * @param revision branch/tag/commit; default `main` + */ + @JvmStatic + @JvmOverloads + public fun fromHuggingFace( + repoId: String, + revision: String = "main", + ctx: ExecutionContext = PooledExecutionContext(DirectCpuExecutionContext()), + mode: BertExecutionMode = BertExecutionMode.DIRECT, + ): EmbeddingModel { + val modelDir = downloadSnapshot(repoId, revision) + return fromSafeTensors(modelDir, ctx, mode, modelId = repoId) + } + + // ------------------------------------------------------------------ + + /** Required snapshot files; download fails loudly when one is missing. */ + private val REQUIRED_FILES = listOf("config.json", "model.safetensors") + + /** Optional snapshot files; absence is fine (bias-free heads, tokenizer variants). */ + private val OPTIONAL_FILES = listOf( + "vocab.txt", + "tokenizer.json", + "2_Dense/config.json", + "2_Dense/model.safetensors", + ) + + private fun downloadSnapshot(repoId: String, revision: String): Path { + val home = System.getProperty("user.home") + ?: error("user.home not set — cannot locate the SKaiNET model cache") + val targetDir = Path.of(home, ".cache", "skainet", "models", repoId.replace('/', '_')) + val missingRequired = REQUIRED_FILES.filterNot { targetDir.resolve(it).exists() } + val missingOptional = OPTIONAL_FILES.filterNot { targetDir.resolve(it).exists() } + if (missingRequired.isEmpty() && missingOptional.isEmpty()) return targetDir + // A previously completed download is detected by the required files + + // tokenizer being present; only then do we skip the network entirely. + if (missingRequired.isEmpty() && targetDir.resolve(DOWNLOAD_MARKER).exists()) return targetDir + + val resolver = JvmDataSourceResolver(useEnvironmentHuggingFaceToken = true) + runBlocking { + missingRequired.forEach { file -> + downloadFile(resolver, repoId, revision, file, targetDir, required = true) + } + missingOptional.forEach { file -> + downloadFile(resolver, repoId, revision, file, targetDir, required = false) + } + } + require(targetDir.resolve("vocab.txt").exists() || targetDir.resolve("tokenizer.json").exists()) { + "Neither vocab.txt nor tokenizer.json available in $repoId — cannot build a tokenizer" + } + Files.write(targetDir.resolve(DOWNLOAD_MARKER), ByteArray(0)) + return targetDir + } + + private suspend fun downloadFile( + resolver: JvmDataSourceResolver, + repoId: String, + revision: String, + file: String, + targetDir: Path, + required: Boolean, + ) { + val target = targetDir.resolve(file) + target.parent.createDirectories() + // Stream to a .part sibling and move atomically so an interrupted + // download retries instead of leaving a truncated file behind. + val tmp = target.resolveSibling(target.name + ".part") + try { + val artifact = resolver.resolve( + DataSourceRequest( + uri = "hf://$repoId@$revision/$file", + // targetDir itself is the cache; the existence check above + // is the cache hit, so skip the resolver's flat store. + cachePolicy = CachePolicy.Bypass, + ) + ) + SystemFileSystem.sink(KotlinxPath(tmp.toString())).buffered().use { sink -> + artifact.copyTo(sink) + } + Files.move(tmp, target, StandardCopyOption.REPLACE_EXISTING) + } catch (e: Exception) { + Files.deleteIfExists(tmp) + if (required) { + throw DataSourceException("Failed to download required file $file from $repoId", e) + } + // Optional file (bias-free 2_Dense heads lack linear.bias etc.) — skip. + } + } + + private const val DOWNLOAD_MARKER = ".skainet-snapshot-complete" + + private fun resolveWeightsFile(modelDir: Path): Path { + listOf("model.safetensors", "pytorch_model.safetensors").forEach { name -> + val p = modelDir.resolve(name) + if (p.exists()) return p + } + modelDir.toFile().listFiles() + ?.firstOrNull { it.extension == "safetensors" } + ?.let { return it.toPath() } + error("No .safetensors file found in $modelDir") + } + + private fun loadTokenizer(modelDir: Path): HuggingFaceTokenizer { + val vocab = modelDir.resolve("vocab.txt") + if (vocab.exists()) return HuggingFaceTokenizer.fromVocabTxt(vocab.readText()) + val tokenizerJson = modelDir.resolve("tokenizer.json") + if (tokenizerJson.exists()) return HuggingFaceTokenizer.fromTokenizerJson(tokenizerJson.readText()) + error("Neither vocab.txt nor tokenizer.json found in $modelDir") + } +} diff --git a/llm-providers/src/main/kotlin/sk/ainet/llm/providers/SkaiNetEmbeddingModel.kt b/llm-providers/src/main/kotlin/sk/ainet/llm/providers/SkaiNetEmbeddingModel.kt index ce4d319f..5dbbd6f2 100644 --- a/llm-providers/src/main/kotlin/sk/ainet/llm/providers/SkaiNetEmbeddingModel.kt +++ b/llm-providers/src/main/kotlin/sk/ainet/llm/providers/SkaiNetEmbeddingModel.kt @@ -1,5 +1,3 @@ -@file:Suppress("DEPRECATION") - package sk.ainet.llm.providers import sk.ainet.apps.llm.Tokenizer @@ -9,21 +7,21 @@ import sk.ainet.llm.api.EmbeddingModel import sk.ainet.llm.api.EmbeddingRequest import sk.ainet.llm.api.EmbeddingResponse import sk.ainet.llm.api.Usage -import sk.ainet.models.bert.BertRuntime +import sk.ainet.models.bert.BertEncoderRuntime /** - * Adapts a BERT-style encoder runtime (`BertRuntime`) + `Tokenizer` to the neutral - * [EmbeddingModel] SPI. + * Adapts a [BertEncoderRuntime] + [Tokenizer] to the neutral [EmbeddingModel] SPI. * - * The runtime already does mean pooling + L2 normalization internally, so the - * adapter is little more than a tensor-to-FloatArray copy. + * The runtime already does mean pooling, the optional sentence-transformers + * projection, and L2 normalization internally, so the adapter is little more + * than tokenize → encode. * * **Threading:** Like [SkaiNetChatModel], a single instance is **not** thread-safe. */ public class SkaiNetEmbeddingModel( - private val runtime: BertRuntime, + private val runtime: BertEncoderRuntime, private val tokenizer: Tokenizer, - public override val dimensions: Int, + public override val dimensions: Int = runtime.dimensions, private val modelId: String? = null, ) : EmbeddingModel { @@ -32,9 +30,7 @@ public class SkaiNetEmbeddingModel( val embeddings = request.inputs.mapIndexed { idx, text -> val tokens = tokenizer.encode(text) totalPromptTokens += tokens.size - val tensor = runtime.encode(tokens) - val vector = tensor.data.copyToFloatArray() - Embedding(index = idx, vector = vector) + Embedding(index = idx, vector = runtime.encode(tokens)) } return EmbeddingResponse( embeddings = embeddings, diff --git a/llm-providers/src/test/kotlin/sk/ainet/llm/providers/BertEmbeddingModelTest.kt b/llm-providers/src/test/kotlin/sk/ainet/llm/providers/BertEmbeddingModelTest.kt new file mode 100644 index 00000000..a125d612 --- /dev/null +++ b/llm-providers/src/test/kotlin/sk/ainet/llm/providers/BertEmbeddingModelTest.kt @@ -0,0 +1,69 @@ +package sk.ainet.llm.providers + +import org.junit.jupiter.api.Assumptions.assumeTrue +import org.junit.jupiter.api.Tag +import org.junit.jupiter.api.Test +import java.nio.file.Path +import kotlin.io.path.isDirectory +import kotlin.math.abs +import kotlin.math.sqrt +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * One-call factory tests on a real LEAF snapshot. Both tests self-skip when + * no model is available; [fromHuggingFace_downloadsAndEmbeds] additionally + * needs network on a cold cache, so it is integration-tagged. + */ +@Tag("integration") +class BertEmbeddingModelTest { + + private fun localSnapshot(): Path? { + System.getenv("LEAF_MODEL_DIR")?.let { p -> + Path.of(p).takeIf { it.isDirectory() }?.let { return it } + } + val home = System.getProperty("user.home") ?: return null + return Path.of(home, ".cache", "skainet", "models", "MongoDB_mdbr-leaf-mt") + .takeIf { it.isDirectory() } + } + + private fun norm(v: FloatArray) = sqrt(v.sumOf { (it * it).toDouble() }).toFloat() + + private fun cosine(a: FloatArray, b: FloatArray): Float { + var dot = 0f + for (i in a.indices) dot += a[i] * b[i] + return dot / (norm(a) * norm(b)) + } + + @Test + fun fromSafeTensors_embedsAndRanksSemantically() { + val dir = localSnapshot() + assumeTrue(dir != null, "No local LEAF snapshot — skipping") + val model = BertEmbeddingModel.fromSafeTensors(dir!!) + + val query = model.embed("How do I reset my password?") + assertEquals(model.dimensions, query.size) + assertTrue(abs(norm(query) - 1f) < 0.01f, "embedding must be L2-normalized") + + val related = model.embed("Steps to change a forgotten password") + val unrelated = model.embed("The quick brown fox jumps over the lazy dog") + val simRelated = cosine(query, related) + val simUnrelated = cosine(query, unrelated) + assertTrue( + simRelated > simUnrelated, + "semantic ranking broken: related=$simRelated unrelated=$simUnrelated" + ) + } + + @Test + fun fromHuggingFace_downloadsAndEmbeds() { + // Uses the shared skainet model cache: instant when warm, downloads + // (~90 MB) when cold. Self-skips only if the machine is fully offline + // AND the cache is cold — surfaced as a DataSourceException then. + val model = BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-mt") + val v = model.embed("hello world") + assertEquals(model.dimensions, v.size) + assertEquals(1024, v.size, "mdbr-leaf-mt projects to 1024 dims") + assertTrue(abs(norm(v) - 1f) < 0.01f) + } +} diff --git a/llm-runtime/kgemma/api/jvm/kgemma.api b/llm-runtime/kgemma/api/jvm/kgemma.api index 955e8e53..6a9519e1 100644 --- a/llm-runtime/kgemma/api/jvm/kgemma.api +++ b/llm-runtime/kgemma/api/jvm/kgemma.api @@ -1,3 +1,59 @@ +public final class sk/ainet/apps/kgemma/FunctionGemma { + public static final field Companion Lsk/ainet/apps/kgemma/FunctionGemma$Companion; + public synthetic fun (Ljava/lang/String;Lsk/ainet/apps/llm/InferenceRuntime;Lsk/ainet/apps/llm/tokenizer/GGUFTokenizer;IIILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun call (Ljava/lang/String;I)Lsk/ainet/apps/kgemma/FunctionGemma$Turn; + public static synthetic fun call$default (Lsk/ainet/apps/kgemma/FunctionGemma;Ljava/lang/String;IILjava/lang/Object;)Lsk/ainet/apps/kgemma/FunctionGemma$Turn; + public final fun exportCompiled (Ljava/lang/String;IZ)Lsk/ainet/apps/kgemma/FunctionGemmaExport$Result; + public static synthetic fun exportCompiled$default (Lsk/ainet/apps/kgemma/FunctionGemma;Ljava/lang/String;IZILjava/lang/Object;)Lsk/ainet/apps/kgemma/FunctionGemmaExport$Result; +} + +public final class sk/ainet/apps/kgemma/FunctionGemma$Companion { + public final fun fromGguf (Ljava/lang/String;F)Lsk/ainet/apps/kgemma/FunctionGemma; + public static synthetic fun fromGguf$default (Lsk/ainet/apps/kgemma/FunctionGemma$Companion;Ljava/lang/String;FILjava/lang/Object;)Lsk/ainet/apps/kgemma/FunctionGemma; +} + +public final class sk/ainet/apps/kgemma/FunctionGemma$Turn { + public fun (Ljava/lang/String;Ljava/util/List;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/util/List; + public final fun copy (Ljava/lang/String;Ljava/util/List;)Lsk/ainet/apps/kgemma/FunctionGemma$Turn; + public static synthetic fun copy$default (Lsk/ainet/apps/kgemma/FunctionGemma$Turn;Ljava/lang/String;Ljava/util/List;ILjava/lang/Object;)Lsk/ainet/apps/kgemma/FunctionGemma$Turn; + public fun equals (Ljava/lang/Object;)Z + public final fun getCalls ()Ljava/util/List; + public final fun getText ()Ljava/lang/String; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/apps/kgemma/FunctionGemmaExport { + public static final field INSTANCE Lsk/ainet/apps/kgemma/FunctionGemmaExport; + public final fun export (Ljava/lang/String;Ljava/lang/String;IFZ)Lsk/ainet/apps/kgemma/FunctionGemmaExport$Result; + public static synthetic fun export$default (Lsk/ainet/apps/kgemma/FunctionGemmaExport;Ljava/lang/String;Ljava/lang/String;IFZILjava/lang/Object;)Lsk/ainet/apps/kgemma/FunctionGemmaExport$Result; +} + +public final class sk/ainet/apps/kgemma/FunctionGemmaExport$Result { + public fun (Ljava/lang/String;Ljava/lang/String;IJI)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/String; + public final fun component3 ()I + public final fun component4 ()J + public final fun component5 ()I + public final fun copy (Ljava/lang/String;Ljava/lang/String;IJI)Lsk/ainet/apps/kgemma/FunctionGemmaExport$Result; + public static synthetic fun copy$default (Lsk/ainet/apps/kgemma/FunctionGemmaExport$Result;Ljava/lang/String;Ljava/lang/String;IJIILjava/lang/Object;)Lsk/ainet/apps/kgemma/FunctionGemmaExport$Result; + public fun equals (Ljava/lang/Object;)Z + public final fun getExternalParamCount ()I + public final fun getMlirPath ()Ljava/lang/String; + public final fun getSafetensorsPath ()Ljava/lang/String; + public final fun getSeq ()I + public final fun getWeightMiB ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/apps/kgemma/FunctionGemmaExportMainKt { + public static final fun main ([Ljava/lang/String;)V +} + public final class sk/ainet/apps/kgemma/Gemma3nIngestion { public fun (Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/apps/kgemma/Gemma3nLoadConfig;)V public synthetic fun (Lsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;Lsk/ainet/apps/kgemma/Gemma3nLoadConfig;ILkotlin/jvm/internal/DefaultConstructorMarker;)V diff --git a/transformer-core/api/jvm/transformer-core.api b/transformer-core/api/jvm/transformer-core.api index 71078a1b..65604c59 100644 --- a/transformer-core/api/jvm/transformer-core.api +++ b/transformer-core/api/jvm/transformer-core.api @@ -121,6 +121,13 @@ public final class sk/ainet/lang/nn/transformer/AppendKVCache : sk/ainet/lang/nn public fun update (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lkotlin/Pair; } +public final class sk/ainet/lang/nn/transformer/AttentionKV { + public fun (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V + public final fun getK ()Lsk/ainet/lang/tensor/Tensor; + public final fun getOutput ()Lsk/ainet/lang/tensor/Tensor; + public final fun getV ()Lsk/ainet/lang/tensor/Tensor; +} + public final class sk/ainet/lang/nn/transformer/GeGLUFFN : sk/ainet/lang/nn/Module, sk/ainet/lang/nn/topology/ModuleParameters { public fun (IILjava/lang/String;Lkotlin/reflect/KClass;)V public synthetic fun (IILjava/lang/String;Lkotlin/reflect/KClass;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -162,6 +169,7 @@ public final class sk/ainet/lang/nn/transformer/MultiHeadAttention : sk/ainet/la public fun (IIIZZZDLjava/lang/Float;ZZLjava/lang/String;Lsk/ainet/lang/nn/transformer/RoPE;Lsk/ainet/lang/nn/transformer/KVCache;Ljava/lang/Integer;Ljava/lang/Integer;Lkotlin/reflect/KClass;)V public synthetic fun (IIIZZZDLjava/lang/Float;ZZLjava/lang/String;Lsk/ainet/lang/nn/transformer/RoPE;Lsk/ainet/lang/nn/transformer/KVCache;Ljava/lang/Integer;Ljava/lang/Integer;Lkotlin/reflect/KClass;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; + public final fun forwardWithKV (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/nn/transformer/AttentionKV; public final fun getAttentionScale ()Ljava/lang/Float; public final fun getBias ()Z public final fun getCausal ()Z @@ -217,6 +225,13 @@ public final class sk/ainet/lang/nn/transformer/PaddedSharedPositionalKVCache : public fun update (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lkotlin/Pair; } +public final class sk/ainet/lang/nn/transformer/PhaseProfile { + public static final field INSTANCE Lsk/ainet/lang/nn/transformer/PhaseProfile; + public final fun report ()Ljava/lang/String; + public final fun reset ()V + public final fun time (Ljava/lang/String;Lkotlin/jvm/functions/Function0;)Ljava/lang/Object; +} + public final class sk/ainet/lang/nn/transformer/PositionalKVCache : sk/ainet/lang/nn/transformer/KVCache { public fun (IIILjava/lang/String;)V public synthetic fun (IIILjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -238,7 +253,9 @@ public final class sk/ainet/lang/nn/transformer/ResidualAdd : sk/ainet/lang/nn/M public final class sk/ainet/lang/nn/transformer/RoPE : sk/ainet/lang/nn/Module { public fun (IIFLsk/ainet/lang/nn/transformer/RoPEMode;Lsk/ainet/lang/nn/transformer/RoPEScaling;FFZLjava/lang/String;)V public synthetic fun (IIFLsk/ainet/lang/nn/transformer/RoPEMode;Lsk/ainet/lang/nn/transformer/RoPEScaling;FFZLjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun buildInterleavedCosSin (II)Lkotlin/Pair; public final fun forward (Lsk/ainet/lang/tensor/Tensor;ILsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; + public final fun forwardWithCosSin (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; public final fun getFreqDenomRotaryDim ()Z public final fun getHeadDim ()I public final fun getMaxSeqLen ()I