Add EmbeddingInversion (vec2text) example: weight-export scripts + RE…

EmbeddingInversion/.gitignore

@@ -0,0 +1,7 @@
+# Converted model weights + goldens are generated offline, not checked in (~1.1 GB+).
+models/
+
+# Gradle
+.gradle/
+build/
+**/build/

EmbeddingInversion/README.md

@@ -0,0 +1,128 @@
+# EmbeddingInversion — vec2text in the SKaiNET ecosystem
+
+Decode sentence embeddings **back into text**, in pure Kotlin. This is a Kotlin port of
+[vec2text](https://github.com/vec2text/vec2text) (Morris et al.) inference, running on
+[SKaiNET](https://github.com/SKaiNET) — the KotlinDL-successor multiplatform DL engine.
+
+Given only the 768-d embedding vector of a sentence (from `sentence-transformers/gtr-t5-base`),
+the models reconstruct an approximation of the original text:
+
+```
+original:      jack morris is a phd student at cornell tech in new york city
+step 0 (invert):  cos 0.71  "jack morris is a neophyte at COL Tech ... New York City University in"
+step 3 (correct): cos 0.83  "jack morris is a ph.D. tech at neorthell alumnus at Cornell University in"
+```
+
+It is a striking privacy demonstration: **embeddings are not anonymous** — text can be
+recovered from them.
+
+## How it works
+
+Three models (all `t5-base` scale), from the `jxm/gtr__nq__32` checkpoints:
+
+1. **GTR embedder** — `sentence-transformers/gtr-t5-base` T5 encoder + mean pooling → target embedding.
+2. **Inversion model** ("hypothesizer", `jxm/gtr__nq__32`) — projects the embedding to 16
+   pseudo-tokens, feeds them to a T5 encoder-decoder, greedily decodes a first guess.
+3. **Corrector** (`jxm/gtr__nq__32__correct`) — given the target, the current hypothesis and their
+   difference, generates a refined guess. Iterated; the hypothesis with the highest cosine
+   similarity to the target wins.
+
+The Kotlin models live upstream in **SKaiNET-transformers**:
+- `llm-inference:t5` — the T5 encoder-decoder runtime + GTR embedder.
+- `llm-inference:vec2text` — inversion model, corrector, and the iterative `Vec2TextInverter`.
+
+## Model weights
+
+SKaiNET loads flat **fp16 SafeTensors**. The reference vec2text checkpoints are custom
+PyTorch classes on HuggingFace, so they need a one-time offline conversion (done once with any
+Python + HuggingFace tooling — this Kotlin repo intentionally ships no conversion scripts).
+Drop the resulting files into `models/` (git-ignored, ~1.1 GB):
+
+| Output file | Source checkpoint | Keep these keys |
+|---|---|---|
+| `gtr_encoder.safetensors` | `sentence-transformers/gtr-t5-base` | `shared.weight`, `encoder.*` |
+| `inversion.safetensors` | `jxm/gtr__nq__32` | `embedding_transform.{0,3}.{weight,bias}`, `encoder_decoder.*` |
+| `corrector.safetensors` | `jxm/gtr__nq__32__correct` | `embedding_transform_{1,2,3}.{0,3}.*`, `layernorm.{weight,bias}`, `encoder_decoder.*` |
+| `tokenizer.json` | `t5-base` | (SentencePiece, shared by all models) |
+
+Conversion notes:
+- Cast float tensors to fp16 and `clone()` them — T5 ties `shared.weight` /
+  `encoder.embed_tokens.weight` / `decoder.embed_tokens.weight` / `lm_head.weight` to one
+  storage, so **keep only `shared.weight`** and drop the aliases (the Kotlin loader feeds it to
+  every tied site; SafeTensors also refuses aliased storage).
+- The `jxm/*` checkpoints are sharded `pytorch_model-0000N-of-…bin` — merge all shards via the
+  `.index.json` before filtering keys.
+- Some older corrector checkpoints share one `embedding_transform.*` MLP; remap it into the three
+  `embedding_transform_{1,2,3}.*` (mirrors `Corrector._remap_state_dict`).
+
+The parity tests (below) compare against golden tensors dumped from the reference models:
+`GtrEmbedderParityTest` needs token ids + the mean-pooled embedding from `gtr-t5-base`;
+`Vec2TextRoundTripTest` needs only the converted weights above.
+
+## Run it
+
+Put the converted weights in `models/`, then either:
+
+```bash
+# Desktop GUI (Round trip + Vector arithmetic tabs):
+./gradlew :app:run
+
+# Or the CLI:
+./gradlew :cli:run --args="jack morris is a phd student at cornell tech in new york city"
+```
+
+CLI environment knobs: `VEC2TEXT_MODELS_DIR` (default `../models` for the app, `./models` for
+the CLI), `VEC2TEXT_STEPS` (default 5), `VEC2TEXT_BEAM` (sequence beam width, default 1 = greedy),
+`VEC2TEXT_TOKEN_BEAMS` (T5 token-level beam, default = `VEC2TEXT_BEAM`). The GUI has a **Beam
+width** slider. Example:
+
+```bash
+VEC2TEXT_BEAM=3 ./gradlew :cli:run --args="jack morris is a phd student at cornell tech in new york city"
+# beam ×3 lifts cosine ~0.77 → ~0.82 vs greedy at one step (and is proportionally slower).
+```
+
+### Build setup
+
+SKaiNET core comes from **Maven Central 0.36.0** (aligned by the `sk.ainet:skainet-bom` platform).
+The inversion models `sk.ainet.transformers:skainet-transformers-inference-{t5,vec2text}` are used
+at **0.37.0** (adds beam search); until that lands on Central it's resolved from the **local Maven
+cache** (a scoped `mavenLocal`). To (re)publish it, from a `SKaiNET-transformers` checkout with
+`VERSION_NAME=0.37.0`:
+
+```bash
+./gradlew :llm-bom:publishToMavenLocal :transformer-core:publishToMavenLocal \
+          :llm-core:publishToMavenLocal :llm-inference:t5:publishToMavenLocal \
+          :llm-inference:vec2text:publishToMavenLocal -PsignAllPublications=false
+```
+
+No composite build, no source checkouts. Once transformers 0.37.0 is on Central, drop the
+`mavenLocal` repository from `cli/` and `app/` `build.gradle.kts`.
+
+> Reconstruction quality scales with correction `steps` and beam width; greedy + few steps + fp16
+> can produce rough or `<unk>`-laden output on short inputs. A decode KV-cache (much faster,
+> compounding with beam) is the remaining M5 follow-up.
+
+## Compose desktop app (`app/`)
+
+`./gradlew :app:run` opens a Compose for Desktop window with two tabs:
+
+- **Round trip** — type text → embed → 768-d embedding strip → invert; the per-step hypotheses
+  stream in with a live cosine bar as each correction completes.
+- **Vector arithmetic** — interpolate two sentence embeddings with a slider and invert the
+  blend (text you never wrote, decoded from a vector — why inversion matters for privacy).
+
+## Status
+
+| Milestone | State |
+|---|---|
+| M0 weight export + golden | ✅ done (scripts here) |
+| M1 T5 encoder + GTR embedder | ✅ verified (cosine 0.99999985 vs reference) |
+| M2 inversion (single-shot) | ✅ working end-to-end |
+| M3 corrector loop | ✅ working end-to-end |
+| M4 runnable CLI | ✅ `./gradlew :cli:run` (Maven Central 0.36.0) |
+| M4 Compose desktop app | ✅ `./gradlew :app:run` — Round trip + Vector arithmetic tabs |
+| M5 beam search | ✅ `VEC2TEXT_BEAM` / GUI slider (transformers 0.37.0) |
+| M5 decode KV-cache speedup | ⏳ follow-up |
+
+Current decoding is greedy with a no-KV-cache O(L²) loop — correct but slow on CPU. Beam search
+and a KV cache (much faster, closer reconstructions) are the main follow-ups.

EmbeddingInversion/app/build.gradle.kts

@@ -0,0 +1,50 @@
+import org.jetbrains.compose.desktop.application.dsl.TargetFormat
+
+plugins {
+    kotlin("jvm")
+    id("org.jetbrains.compose")
+    id("org.jetbrains.kotlin.plugin.compose")
+}
+
+repositories {
+    google()
+    mavenCentral()
+    // SKaiNET-transformers 0.37.0 (beam search) from the local Maven cache until it lands on Central.
+    mavenLocal { mavenContent { includeGroupAndSubgroups("sk.ainet") } }
+}
+
+kotlin {
+    jvmToolchain(21)
+}
+
+dependencies {
+    // SKaiNET core from Maven Central, version aligned by the published BOM (0.36.0).
+    implementation(platform("sk.ainet:skainet-bom:0.36.0"))
+    implementation("sk.ainet.core:skainet-lang-core")
+    implementation("sk.ainet.core:skainet-backend-cpu")
+    implementation("sk.ainet.core:skainet-io-core")
+    implementation("sk.ainet.core:skainet-io-safetensors")
+
+    // t5 / vec2text 0.37.0 (beam search) — from the local Maven cache until it lands on Central
+    // (publish with: SKaiNET-transformers `./gradlew publishToMavenLocal -PsignAllPublications=false`).
+    implementation("sk.ainet.transformers:skainet-transformers-inference-t5:0.37.0")
+    implementation("sk.ainet.transformers:skainet-transformers-inference-vec2text:0.37.0")
+
+    implementation("org.jetbrains.kotlinx:kotlinx-coroutines-swing:1.11.0")
+    implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:1.11.0")
+
+    implementation(compose.desktop.currentOs)
+    implementation(compose.material3)
+}
+
+compose.desktop {
+    application {
+        mainClass = "sk.ainet.samples.vec2text.ui.MainKt"
+        jvmArgs += listOf("-Xmx4g")
+        nativeDistributions {
+            targetFormats(TargetFormat.Dmg, TargetFormat.Msi, TargetFormat.Deb)
+            packageName = "EmbeddingInversion"
+            packageVersion = "1.0.0"
+        }
+    }
+}

EmbeddingInversion/app/src/main/kotlin/sk/ainet/samples/vec2text/ui/Engine.kt

@@ -0,0 +1,156 @@
+package sk.ainet.samples.vec2text.ui
+
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.jsonObject
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.io.JvmRandomAccessSource
+import sk.ainet.io.safetensors.SafeTensorsParametersLoader
+import sk.ainet.io.tokenizer.SentencePieceTokenizer
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.tensor.plus
+import sk.ainet.lang.tensor.times
+import sk.ainet.lang.types.FP32
+import sk.ainet.models.t5.GtrEmbedder
+import sk.ainet.models.t5.T5Config
+import sk.ainet.models.t5.T5Runtime
+import sk.ainet.models.t5.loadT5Weights
+import sk.ainet.models.vec2text.CorrectorModel
+import sk.ainet.models.vec2text.InversionModel
+import sk.ainet.models.vec2text.Vec2TextInverter
+import sk.ainet.models.vec2text.Vec2TextWeightLoader
+import java.io.File
+
+/** One correction step surfaced to the UI. */
+data class Step(val step: Int, val text: String, val cosine: Float)
+
+/**
+ * Loads the gtr-base checkpoints once and drives the vec2text inversion loop **step by step**
+ * (using the public InversionModel / CorrectorModel APIs) so the UI can render each hypothesis
+ * as it's produced — the full loop takes minutes on CPU.
+ */
+class Vec2TextEngine private constructor(
+    private val embedder: GtrEmbedder<FP32>,
+    private val inversion: InversionModel<FP32>,
+    private val corrector: CorrectorModel<FP32>,
+    private val sp: SentencePieceTokenizer,
+    private val cfg: T5Config,
+) {
+    private fun encodeForEmbedder(text: String): IntArray {
+        val ids = sp.encode(text).take(cfg.maxSeqLength - 1).toMutableList()
+        ids.add(cfg.eosTokenId)
+        return ids.toIntArray()
+    }
+
+    private fun decode(ids: IntArray): String =
+        sp.decode(ids.filter { it != 0 && it != cfg.eosTokenId }.toIntArray())
+
+    /** Embed [text] into its `[768]` GTR sentence embedding. */
+    fun embed(text: String): Tensor<FP32, Float> = embedder.embed(encodeForEmbedder(text))
+
+    /** Raw float view of an embedding, for visualization. */
+    fun toFloats(v: Tensor<FP32, Float>): FloatArray = v.data.copyToFloatArray()
+
+    /** Linear interpolation `(1-alpha)*a + alpha*b` of two embeddings. */
+    fun interpolate(a: Tensor<FP32, Float>, b: Tensor<FP32, Float>, alpha: Float): Tensor<FP32, Float> =
+        (a * (1f - alpha)) + (b * alpha)
+
+    /**
+     * Invert [target] into text over [steps] correction rounds, calling [onStep] after the
+     * initial hypothesis (step 0) and each correction. Returns the best-cosine hypothesis.
+     *
+     * [beamWidth] > 1 or [tokenBeams] > 1 enables beam search (better reconstructions, slower):
+     * [tokenBeams] is the T5 token-level beam per generation; [beamWidth] keeps that many
+     * hypotheses across correction rounds, ranked by cosine to [target].
+     */
+    fun invert(
+        target: Tensor<FP32, Float>,
+        steps: Int,
+        beamWidth: Int = 1,
+        tokenBeams: Int = 1,
+        onStep: (Step) -> Unit,
+    ): Step =
+        if (beamWidth <= 1 && tokenBeams <= 1) invertGreedy(target, steps, onStep)
+        else invertBeam(target, steps, beamWidth.coerceAtLeast(1), tokenBeams.coerceAtLeast(1), onStep)
+
+    private fun invertGreedy(target: Tensor<FP32, Float>, steps: Int, onStep: (Step) -> Unit): Step {
+        var hypIds = inversion.invert(target, maxLength = cfg.maxSeqLength)
+        var hypText = decode(hypIds)
+        var best = Step(0, hypText, cosineOf(target, hypText))
+        onStep(best)
+
+        for (s in 1..steps) {
+            hypIds = corrector.correct(target, embed(hypText), hypIds, maxLength = cfg.maxSeqLength)
+            hypText = decode(hypIds)
+            val step = Step(s, hypText, cosineOf(target, hypText))
+            onStep(step)
+            if (step.cosine > best.cosine) best = step
+        }
+        return best
+    }
+
+    /** Streaming sequence-level beam: emit the best-of-beam hypothesis after each round. */
+    private fun invertBeam(target: Tensor<FP32, Float>, steps: Int, beamWidth: Int, tokenBeams: Int, onStep: (Step) -> Unit): Step {
+        var beams = rank(target, inversion.invertBeam(target, maxOf(beamWidth, tokenBeams), cfg.maxSeqLength)).take(beamWidth)
+        var best = Step(0, beams.first().text, beams.first().cos)
+        onStep(best)
+
+        for (s in 1..steps) {
+            val pool = ArrayList<IntArray>()
+            for (b in beams) pool += corrector.correctBeam(target, embed(b.text), b.ids, tokenBeams, cfg.maxSeqLength)
+            beams = rank(target, pool).take(beamWidth)
+            val sb = beams.first()
+            val step = Step(s, sb.text, sb.cos)
+            onStep(step)
+            if (step.cosine > best.cosine) best = step
+        }
+        return best
+    }
+
+    private class Cand(val ids: IntArray, val text: String, val cos: Float)
+
+    private fun rank(target: Tensor<FP32, Float>, idsList: List<IntArray>): List<Cand> =
+        idsList.asSequence()
+            .map { ids -> decode(ids) to ids }
+            .distinctBy { it.first }
+            .map { (text, ids) -> Cand(ids, text, cosineOf(target, text)) }
+            .sortedByDescending { it.cos }
+            .toList()
+
+    private fun cosineOf(target: Tensor<FP32, Float>, text: String): Float =
+        Vec2TextInverter.cosine(target, embed(text))
+
+    companion object {
+        val REQUIRED = listOf(
+            "tokenizer.json", "gtr_encoder.safetensors", "inversion.safetensors", "corrector.safetensors",
+        )
+
+        /** Resolve the models dir: `$VEC2TEXT_MODELS_DIR`, else `../models`, else `models`. */
+        fun modelsDir(): File {
+            System.getenv("VEC2TEXT_MODELS_DIR")?.let { return File(it) }
+            for (c in listOf("../models", "models")) {
+                val f = File(c)
+                if (File(f, "tokenizer.json").exists()) return f
+            }
+            return File("../models")
+        }
+
+        fun missing(dir: File): List<String> = REQUIRED.filterNot { File(dir, it).exists() }
+
+        /** Blocking load of all three models (call off the UI thread). */
+        suspend fun load(dir: File): Vec2TextEngine {
+            val ctx = DirectCpuExecutionContext()
+            val cfg = T5Config()
+            fun loader(name: String) =
+                SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(File(dir, name).toString()) })
+
+            val gtr = loadT5Weights(loader("gtr_encoder.safetensors"), ctx, FP32::class, cfg, "", withDecoder = false)
+            val embedder = GtrEmbedder(T5Runtime(ctx, gtr, FP32::class))
+            val inversion = InversionModel(ctx, Vec2TextWeightLoader.loadInversion(loader("inversion.safetensors"), ctx, FP32::class, cfg), FP32::class)
+            val corrector = CorrectorModel(ctx, Vec2TextWeightLoader.loadCorrector(loader("corrector.safetensors"), ctx, FP32::class, cfg), FP32::class)
+            val sp = SentencePieceTokenizer.fromTokenizerJson(
+                Json.parseToJsonElement(File(dir, "tokenizer.json").readText()).jsonObject
+            )
+            return Vec2TextEngine(embedder, inversion, corrector, sp, cfg)
+        }
+    }
+}