diff --git a/EmbeddingInversion/.gitignore b/EmbeddingInversion/.gitignore new file mode 100644 index 0000000..dc06c9a --- /dev/null +++ b/EmbeddingInversion/.gitignore @@ -0,0 +1,7 @@ +# Converted model weights + goldens are generated offline, not checked in (~1.1 GB+). +models/ + +# Gradle +.gradle/ +build/ +**/build/ diff --git a/EmbeddingInversion/README.md b/EmbeddingInversion/README.md new file mode 100644 index 0000000..b174d4f --- /dev/null +++ b/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 ``-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. diff --git a/EmbeddingInversion/app/build.gradle.kts b/EmbeddingInversion/app/build.gradle.kts new file mode 100644 index 0000000..140d0a9 --- /dev/null +++ b/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" + } + } +} diff --git a/EmbeddingInversion/app/src/main/kotlin/sk/ainet/samples/vec2text/ui/Engine.kt b/EmbeddingInversion/app/src/main/kotlin/sk/ainet/samples/vec2text/ui/Engine.kt new file mode 100644 index 0000000..f05b984 --- /dev/null +++ b/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, + private val inversion: InversionModel, + private val corrector: CorrectorModel, + 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 = embedder.embed(encodeForEmbedder(text)) + + /** Raw float view of an embedding, for visualization. */ + fun toFloats(v: Tensor): FloatArray = v.data.copyToFloatArray() + + /** Linear interpolation `(1-alpha)*a + alpha*b` of two embeddings. */ + fun interpolate(a: Tensor, b: Tensor, alpha: Float): Tensor = + (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, + 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, 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, 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() + 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, idsList: List): List = + 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, 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 = 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) + } + } +} diff --git a/EmbeddingInversion/app/src/main/kotlin/sk/ainet/samples/vec2text/ui/Main.kt b/EmbeddingInversion/app/src/main/kotlin/sk/ainet/samples/vec2text/ui/Main.kt new file mode 100644 index 0000000..6eef2ba --- /dev/null +++ b/EmbeddingInversion/app/src/main/kotlin/sk/ainet/samples/vec2text/ui/Main.kt @@ -0,0 +1,251 @@ +package sk.ainet.samples.vec2text.ui + +import androidx.compose.foundation.Canvas +import androidx.compose.foundation.layout.Arrangement +import androidx.compose.foundation.layout.Column +import androidx.compose.foundation.layout.Row +import androidx.compose.foundation.layout.Spacer +import androidx.compose.foundation.layout.fillMaxSize +import androidx.compose.foundation.layout.fillMaxWidth +import androidx.compose.foundation.layout.height +import androidx.compose.foundation.layout.padding +import androidx.compose.foundation.rememberScrollState +import androidx.compose.foundation.verticalScroll +import androidx.compose.material3.Button +import androidx.compose.material3.CircularProgressIndicator +import androidx.compose.material3.MaterialTheme +import androidx.compose.material3.OutlinedTextField +import androidx.compose.material3.Slider +import androidx.compose.material3.Tab +import androidx.compose.material3.TabRow +import androidx.compose.material3.Text +import androidx.compose.runtime.Composable +import androidx.compose.runtime.LaunchedEffect +import androidx.compose.runtime.getValue +import androidx.compose.runtime.mutableStateListOf +import androidx.compose.runtime.mutableStateOf +import androidx.compose.runtime.remember +import androidx.compose.runtime.rememberCoroutineScope +import androidx.compose.runtime.setValue +import androidx.compose.ui.Alignment +import androidx.compose.ui.Modifier +import androidx.compose.ui.geometry.Offset +import androidx.compose.ui.graphics.Color +import androidx.compose.ui.unit.dp +import androidx.compose.ui.window.Window +import androidx.compose.ui.window.application +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.launch +import kotlinx.coroutines.withContext +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.math.min + +private sealed interface Load { + data object Loading : Load + data class Missing(val dir: String, val files: List) : Load + data class Ready(val engine: Vec2TextEngine) : Load + data class Failed(val message: String) : Load +} + +fun main() = application { + Window(onCloseRequest = ::exitApplication, title = "Embedding Inversion — vec2text on SKaiNET") { + MaterialTheme { App() } + } +} + +@Composable +private fun App() { + var load by remember { mutableStateOf(Load.Loading) } + LaunchedEffect(Unit) { + val dir = Vec2TextEngine.modelsDir() + val missing = Vec2TextEngine.missing(dir) + load = if (missing.isNotEmpty()) { + Load.Missing(dir.absolutePath, missing) + } else try { + Load.Ready(withContext(Dispatchers.Default) { Vec2TextEngine.load(dir) }) + } catch (e: Throwable) { + Load.Failed(e.message ?: e.toString()) + } + } + + Column(Modifier.fillMaxSize().padding(20.dp)) { + Text("Embedding Inversion", style = MaterialTheme.typography.headlineSmall) + Text( + "Decode a sentence embedding back into text (vec2text on SKaiNET). Embeddings are not anonymous.", + style = MaterialTheme.typography.bodyMedium, + ) + Spacer(Modifier.height(16.dp)) + when (val l = load) { + is Load.Loading -> Row(verticalAlignment = Alignment.CenterVertically) { + CircularProgressIndicator(Modifier.height(20.dp)); Spacer(Modifier.height(8.dp)) + Text(" Loading models (~1.1 GB)…") + } + is Load.Missing -> Text( + "Models not found in ${l.dir}\nMissing: ${l.files.joinToString()}\n" + + "Put the converted weights there or set VEC2TEXT_MODELS_DIR (see README).", + color = MaterialTheme.colorScheme.error, + ) + is Load.Failed -> Text("Failed to load models: ${l.message}", color = MaterialTheme.colorScheme.error) + is Load.Ready -> Tabs(l.engine) + } + } +} + +@Composable +private fun Tabs(engine: Vec2TextEngine) { + var tab by remember { mutableStateOf(0) } + val titles = listOf("Round trip", "Vector arithmetic") + TabRow(selectedTabIndex = tab) { + titles.forEachIndexed { i, t -> Tab(selected = tab == i, onClick = { tab = i }, text = { Text(t) }) } + } + Spacer(Modifier.height(16.dp)) + when (tab) { + 0 -> RoundTripTab(engine) + else -> VectorArithmeticTab(engine) + } +} + +@Composable +private fun RoundTripTab(engine: Vec2TextEngine) { + val scope = rememberCoroutineScope() + var text by remember { mutableStateOf("jack morris is a phd student at cornell tech in new york city") } + var steps by remember { mutableStateOf(5f) } + var beam by remember { mutableStateOf(1f) } + var running by remember { mutableStateOf(false) } + var vector by remember { mutableStateOf(null) } + val trace = remember { mutableStateListOf() } + + Column(Modifier.verticalScroll(rememberScrollState())) { + OutlinedTextField(text, { text = it }, Modifier.fillMaxWidth(), label = { Text("Text to embed & invert") }) + StepsSlider(steps) { steps = it } + BeamSlider(beam) { beam = it } + RunButton("Embed → invert", running) { + running = true; trace.clear(); vector = null + scope.launch { + withContext(Dispatchers.Default) { + val target = engine.embed(text) + vector = engine.toFloats(target) + engine.invert(target, steps.toInt(), beam.toInt(), beam.toInt()) { trace.add(it) } + } + running = false + } + } + vector?.let { EmbeddingStrip(it) } + Reconstruction(original = text, trace = trace) + } +} + +@Composable +private fun VectorArithmeticTab(engine: Vec2TextEngine) { + val scope = rememberCoroutineScope() + var a by remember { mutableStateOf("the weather in paris is cold and rainy") } + var b by remember { mutableStateOf("the food in tokyo is fresh and delicious") } + var alpha by remember { mutableStateOf(0.5f) } + var running by remember { mutableStateOf(false) } + var vector by remember { mutableStateOf(null) } + val trace = remember { mutableStateListOf() } + + Column(Modifier.verticalScroll(rememberScrollState())) { + OutlinedTextField(a, { a = it }, Modifier.fillMaxWidth(), label = { Text("Sentence A") }) + OutlinedTextField(b, { b = it }, Modifier.fillMaxWidth(), label = { Text("Sentence B") }) + Text("Interpolation A ${"%.2f".format(1 - alpha)} ↔ B ${"%.2f".format(alpha)}") + Slider(alpha, { alpha = it }, valueRange = 0f..1f) + RunButton("Invert interpolated embedding", running) { + running = true; trace.clear(); vector = null + scope.launch { + withContext(Dispatchers.Default) { + val target = engine.interpolate(engine.embed(a), engine.embed(b), alpha) + vector = engine.toFloats(target) + engine.invert(target, 5) { trace.add(it) } + } + running = false + } + } + Text( + "Inverting a blend of two sentence embeddings — text you never wrote, decoded from a vector.", + style = MaterialTheme.typography.bodySmall, + ) + vector?.let { EmbeddingStrip(it) } + Reconstruction(original = "(interpolated embedding)", trace = trace) + } +} + +@Composable +private fun StepsSlider(steps: Float, onChange: (Float) -> Unit) { + Text("Correction steps: ${steps.toInt()}") + Slider(steps, onChange, valueRange = 0f..20f, steps = 19) +} + +@Composable +private fun BeamSlider(beam: Float, onChange: (Float) -> Unit) { + val w = beam.toInt() + Text("Beam width: ${if (w <= 1) "1 (greedy)" else "$w (slower, better)"}") + Slider(beam, onChange, valueRange = 1f..4f, steps = 2) +} + +@Composable +private fun RunButton(label: String, running: Boolean, onClick: () -> Unit) { + Row(verticalAlignment = Alignment.CenterVertically) { + Button(onClick = onClick, enabled = !running) { Text(label) } + if (running) { + Spacer(Modifier.height(8.dp)) + CircularProgressIndicator(Modifier.height(18.dp).padding(start = 12.dp)) + Text(" working… (CPU decode is slow)") + } + } + Spacer(Modifier.height(12.dp)) +} + +/** A horizontal strip visualizing the 768-d embedding: one column per (downsampled) dim, red<0= 0) Color(0f, 0.35f, 1f, value.coerceIn(0f, 1f)) + else Color(1f, 0.2f, 0.2f, (-value).coerceIn(0f, 1f)) + drawRect(c, topLeft = Offset(i * cw, 0f), size = androidx.compose.ui.geometry.Size(cw + 1, size.height)) + } + } + Spacer(Modifier.height(12.dp)) +} + +@Composable +private fun Reconstruction(original: String, trace: List) { + if (trace.isEmpty()) return + Spacer(Modifier.height(12.dp)) + Text("original: $original", style = MaterialTheme.typography.bodyMedium) + val best = trace.maxByOrNull { it.cosine } + if (best != null) { + Text("reconstructed: ${best.text}", style = MaterialTheme.typography.titleMedium) + Text("best cosine: ${"%.4f".format(best.cosine)}") + } + Spacer(Modifier.height(8.dp)) + Text("steps", style = MaterialTheme.typography.labelMedium) + Column { + for (s in trace) { + Row(Modifier.fillMaxWidth().padding(vertical = 2.dp), verticalAlignment = Alignment.CenterVertically) { + Text("${s.step}", Modifier.height(20.dp)) + CosineBar(s.cosine) + Text(" ${"%.3f".format(s.cosine)} ${s.text}", style = MaterialTheme.typography.bodySmall) + } + } + } +} + +@Composable +private fun CosineBar(cosine: Float) { + Canvas(Modifier.height(12.dp).fillMaxWidth(0.18f).padding(horizontal = 8.dp)) { + drawRect(Color(0.85f, 0.85f, 0.85f), size = size) + drawRect( + Color(0.2f, 0.6f, 0.3f), + size = androidx.compose.ui.geometry.Size(size.width * cosine.coerceIn(0f, 1f), size.height), + ) + } +} diff --git a/EmbeddingInversion/build.gradle.kts b/EmbeddingInversion/build.gradle.kts new file mode 100644 index 0000000..4e6287c --- /dev/null +++ b/EmbeddingInversion/build.gradle.kts @@ -0,0 +1,7 @@ +// Root build for the EmbeddingInversion example. Plugin versions are declared once here +// (apply false) so the :cli and :app subprojects can apply them without repeating versions. +plugins { + kotlin("jvm") version "2.4.0" apply false + id("org.jetbrains.compose") version "1.10.1" apply false + id("org.jetbrains.kotlin.plugin.compose") version "2.4.0" apply false +} diff --git a/EmbeddingInversion/cli/build.gradle.kts b/EmbeddingInversion/cli/build.gradle.kts new file mode 100644 index 0000000..d140211 --- /dev/null +++ b/EmbeddingInversion/cli/build.gradle.kts @@ -0,0 +1,36 @@ +plugins { + kotlin("jvm") + application +} + +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: version aligned by the published BOM (0.36.0), resolved from Maven Central. + 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-core:1.11.0") + implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:1.11.0") +} + +application { + mainClass.set("sk.ainet.samples.vec2text.MainKt") +} diff --git a/EmbeddingInversion/cli/src/main/kotlin/sk/ainet/samples/vec2text/Main.kt b/EmbeddingInversion/cli/src/main/kotlin/sk/ainet/samples/vec2text/Main.kt new file mode 100644 index 0000000..3c0c726 --- /dev/null +++ b/EmbeddingInversion/cli/src/main/kotlin/sk/ainet/samples/vec2text/Main.kt @@ -0,0 +1,81 @@ +package sk.ainet.samples.vec2text + +import kotlinx.coroutines.runBlocking +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.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.Vec2TextTokenizer +import sk.ainet.models.vec2text.Vec2TextWeightLoader +import sk.ainet.lang.types.FP32 +import java.io.File + +/** + * EmbeddingInversion CLI — decode a sentence embedding back into text with vec2text on SKaiNET. + * + * Usage: + * ./gradlew :cli:run --args="" + * Model directory (default ./models) can be overridden with VEC2TEXT_MODELS_DIR. + * Needs gtr_encoder / inversion / corrector .safetensors + tokenizer.json (see README). + */ +fun main(args: Array) = runBlocking { + val modelsDir = File(System.getenv("VEC2TEXT_MODELS_DIR") ?: "models") + val text = args.joinToString(" ").ifBlank { + "jack morris is a phd student at cornell tech in new york city" + } + val steps = System.getenv("VEC2TEXT_STEPS")?.toIntOrNull() ?: 5 + val beamWidth = System.getenv("VEC2TEXT_BEAM")?.toIntOrNull() ?: 1 + val tokenBeams = System.getenv("VEC2TEXT_TOKEN_BEAMS")?.toIntOrNull() ?: beamWidth + + val required = listOf("tokenizer.json", "gtr_encoder.safetensors", "inversion.safetensors", "corrector.safetensors") + val missing = required.filterNot { File(modelsDir, it).exists() } + if (missing.isNotEmpty()) { + System.err.println("Missing model files in ${modelsDir.absolutePath}: $missing") + System.err.println("See README.md for how to produce them, or set VEC2TEXT_MODELS_DIR.") + return@runBlocking + } + + val ctx = DirectCpuExecutionContext() + val cfg = T5Config() + fun loader(name: String) = + SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(File(modelsDir, name).toString()) }) + + println("Loading models from ${modelsDir.absolutePath} …") + 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(modelsDir, "tokenizer.json").readText()).jsonObject + ) + val codec = object : Vec2TextTokenizer { + override fun encodeForEmbedder(text: String): IntArray { + val ids = sp.encode(text).take(cfg.maxSeqLength - 1).toMutableList() + ids.add(cfg.eosTokenId) + return ids.toIntArray() + } + override fun decode(ids: IntArray): String = + sp.decode(ids.filter { it != 0 && it != cfg.eosTokenId }.toIntArray()) + } + + val mode = if (beamWidth > 1 || tokenBeams > 1) "beam ×$beamWidth, token-beams ×$tokenBeams" else "greedy" + println("Inverting (≤$steps correction steps, $mode)…\n") + val result = Vec2TextInverter(embedder, inversion, corrector, codec) + .invert(text, numSteps = steps, maxLength = cfg.maxSeqLength, sequenceBeamWidth = beamWidth, tokenBeams = tokenBeams) + + println("original: $text") + println("reconstructed: ${result.text}") + println("cosine: ${"%.4f".format(result.cosine)}\n") + println("trace:") + result.trace.forEach { println(" step ${it.step}: cos=${"%.4f".format(it.cosine)} \"${it.text}\"") } +} diff --git a/EmbeddingInversion/gradle.properties b/EmbeddingInversion/gradle.properties new file mode 100644 index 0000000..6b74e16 --- /dev/null +++ b/EmbeddingInversion/gradle.properties @@ -0,0 +1,4 @@ +kotlin.code.style=official +org.gradle.jvmargs=-Xmx4g -Dfile.encoding=UTF-8 +org.gradle.configuration-cache=true +kotlin.native.ignoreDisabledTargets=true diff --git a/EmbeddingInversion/gradle/wrapper/gradle-wrapper.jar b/EmbeddingInversion/gradle/wrapper/gradle-wrapper.jar new file mode 100644 index 0000000..b1b8ef5 Binary files /dev/null and b/EmbeddingInversion/gradle/wrapper/gradle-wrapper.jar differ diff --git a/EmbeddingInversion/gradle/wrapper/gradle-wrapper.properties b/EmbeddingInversion/gradle/wrapper/gradle-wrapper.properties new file mode 100644 index 0000000..db6d194 --- /dev/null +++ b/EmbeddingInversion/gradle/wrapper/gradle-wrapper.properties @@ -0,0 +1,8 @@ +distributionBase=GRADLE_USER_HOME +distributionPath=wrapper/dists +distributionUrl=https\://services.gradle.org/distributions/gradle-9.6.1-bin.zip +networkTimeout=10000 +retries=0 +retryBackOffMs=500 +zipStoreBase=GRADLE_USER_HOME +zipStorePath=wrapper/dists diff --git a/EmbeddingInversion/gradlew b/EmbeddingInversion/gradlew new file mode 100755 index 0000000..249efbb --- /dev/null +++ b/EmbeddingInversion/gradlew @@ -0,0 +1,248 @@ +#!/bin/sh + +# +# Copyright © 2015 the original authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 +# + +############################################################################## +# +# gradlew start up script for POSIX generated by Gradle. +# +# Important for running: +# +# (1) You need a POSIX-compliant shell to run this script. If your /bin/sh is +# noncompliant, but you have some other compliant shell such as ksh or +# bash, then to run this script, type that shell name before the whole +# command line, like: +# +# ksh gradlew +# +# Busybox and similar reduced shells will NOT work, because this script +# requires all of these POSIX shell features: +# * functions; +# * expansions «$var», «${var}», «${var:-default}», «${var+SET}», +# «${var#prefix}», «${var%suffix}», and «$( cmd )»; +# * compound commands having a testable exit status, especially «case»; +# * various built-in commands including «command», «set», and «ulimit». +# +# Important for patching: +# +# (2) This script targets any POSIX shell, so it avoids extensions provided +# by Bash, Ksh, etc; in particular arrays are avoided. +# +# The "traditional" practice of packing multiple parameters into a +# space-separated string is a well documented source of bugs and security +# problems, so this is (mostly) avoided, by progressively accumulating +# options in "$@", and eventually passing that to Java. +# +# Where the inherited environment variables (DEFAULT_JVM_OPTS, JAVA_OPTS, +# and GRADLE_OPTS) rely on word-splitting, this is performed explicitly; +# see the in-line comments for details. +# +# There are tweaks for specific operating systems such as AIX, CygWin, +# Darwin, MinGW, and NonStop. +# +# (3) This script is generated from the Groovy template +# https://github.com/gradle/gradle/blob/3d91ce3b8caaf77ad09f381f43615b715b53f72c/platforms/jvm/plugins-application/src/main/resources/org/gradle/api/internal/plugins/unixStartScript.txt +# within the Gradle project. +# +# You can find Gradle at https://github.com/gradle/gradle/. +# +############################################################################## + +# Attempt to set APP_HOME + +# Resolve links: $0 may be a link +app_path=$0 + +# Need this for daisy-chained symlinks. +while + APP_HOME=${app_path%"${app_path##*/}"} # leaves a trailing /; empty if no leading path + [ -h "$app_path" ] +do + ls=$( ls -ld "$app_path" ) + link=${ls#*' -> '} + case $link in #( + /*) app_path=$link ;; #( + *) app_path=$APP_HOME$link ;; + esac +done + +# This is normally unused +# shellcheck disable=SC2034 +APP_BASE_NAME=${0##*/} +# Discard cd standard output in case $CDPATH is set (https://github.com/gradle/gradle/issues/25036) +APP_HOME=$( cd -P "${APP_HOME:-./}" > /dev/null && printf '%s\n' "$PWD" ) || exit + +# Use the maximum available, or set MAX_FD != -1 to use that value. +MAX_FD=maximum + +warn () { + echo "$*" +} >&2 + +die () { + echo + echo "$*" + echo + exit 1 +} >&2 + +# OS specific support (must be 'true' or 'false'). +cygwin=false +msys=false +darwin=false +nonstop=false +case "$( uname )" in #( + CYGWIN* ) cygwin=true ;; #( + Darwin* ) darwin=true ;; #( + MSYS* | MINGW* ) msys=true ;; #( + NONSTOP* ) nonstop=true ;; +esac + + + +# Determine the Java command to use to start the JVM. +if [ -n "$JAVA_HOME" ] ; then + if [ -x "$JAVA_HOME/jre/sh/java" ] ; then + # IBM's JDK on AIX uses strange locations for the executables + JAVACMD=$JAVA_HOME/jre/sh/java + else + JAVACMD=$JAVA_HOME/bin/java + fi + if [ ! -x "$JAVACMD" ] ; then + die "ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME + +Please set the JAVA_HOME variable in your environment to match the +location of your Java installation." + fi +else + JAVACMD=java + if ! command -v java >/dev/null 2>&1 + then + die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. + +Please set the JAVA_HOME variable in your environment to match the +location of your Java installation." + fi +fi + +# Increase the maximum file descriptors if we can. +if ! "$cygwin" && ! "$darwin" && ! "$nonstop" ; then + case $MAX_FD in #( + max*) + # In POSIX sh, ulimit -H is undefined. That's why the result is checked to see if it worked. + # shellcheck disable=SC2039,SC3045 + MAX_FD=$( ulimit -H -n ) || + warn "Could not query maximum file descriptor limit" + esac + case $MAX_FD in #( + '' | soft) :;; #( + *) + # In POSIX sh, ulimit -n is undefined. That's why the result is checked to see if it worked. + # shellcheck disable=SC2039,SC3045 + ulimit -n "$MAX_FD" || + warn "Could not set maximum file descriptor limit to $MAX_FD" + esac +fi + +# Collect all arguments for the java command, stacking in reverse order: +# * args from the command line +# * the main class name +# * -classpath +# * -D...appname settings +# * --module-path (only if needed) +# * DEFAULT_JVM_OPTS, JAVA_OPTS, and GRADLE_OPTS environment variables. + +# For Cygwin or MSYS, switch paths to Windows format before running java +if "$cygwin" || "$msys" ; then + APP_HOME=$( cygpath --path --mixed "$APP_HOME" ) + + JAVACMD=$( cygpath --unix "$JAVACMD" ) + + # Now convert the arguments - kludge to limit ourselves to /bin/sh + for arg do + if + case $arg in #( + -*) false ;; # don't mess with options #( + /?*) t=${arg#/} t=/${t%%/*} # looks like a POSIX filepath + [ -e "$t" ] ;; #( + *) false ;; + esac + then + arg=$( cygpath --path --ignore --mixed "$arg" ) + fi + # Roll the args list around exactly as many times as the number of + # args, so each arg winds up back in the position where it started, but + # possibly modified. + # + # NB: a `for` loop captures its iteration list before it begins, so + # changing the positional parameters here affects neither the number of + # iterations, nor the values presented in `arg`. + shift # remove old arg + set -- "$@" "$arg" # push replacement arg + done +fi + + +# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. +DEFAULT_JVM_OPTS='"-Xmx64m" "-Xms64m"' + +# Collect all arguments for the java command: +# * DEFAULT_JVM_OPTS, JAVA_OPTS, and optsEnvironmentVar are not allowed to contain shell fragments, +# and any embedded shellness will be escaped. +# * For example: A user cannot expect ${Hostname} to be expanded, as it is an environment variable and will be +# treated as '${Hostname}' itself on the command line. + +set -- \ + "-Dorg.gradle.appname=$APP_BASE_NAME" \ + -jar "$APP_HOME/gradle/wrapper/gradle-wrapper.jar" \ + "$@" + +# Stop when "xargs" is not available. +if ! command -v xargs >/dev/null 2>&1 +then + die "xargs is not available" +fi + +# Use "xargs" to parse quoted args. +# +# With -n1 it outputs one arg per line, with the quotes and backslashes removed. +# +# In Bash we could simply go: +# +# readarray ARGS < <( xargs -n1 <<<"$var" ) && +# set -- "${ARGS[@]}" "$@" +# +# but POSIX shell has neither arrays nor command substitution, so instead we +# post-process each arg (as a line of input to sed) to backslash-escape any +# character that might be a shell metacharacter, then use eval to reverse +# that process (while maintaining the separation between arguments), and wrap +# the whole thing up as a single "set" statement. +# +# This will of course break if any of these variables contains a newline or +# an unmatched quote. +# + +eval "set -- $( + printf '%s\n' "$DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS" | + xargs -n1 | + sed ' s~[^-[:alnum:]+,./:=@_]~\\&~g; ' | + tr '\n' ' ' + )" '"$@"' + +exec "$JAVACMD" "$@" diff --git a/EmbeddingInversion/gradlew.bat b/EmbeddingInversion/gradlew.bat new file mode 100644 index 0000000..8508ef6 --- /dev/null +++ b/EmbeddingInversion/gradlew.bat @@ -0,0 +1,82 @@ +@rem +@rem Copyright 2015 the original author or authors. +@rem +@rem Licensed under the Apache License, Version 2.0 (the "License"); +@rem you may not use this file except in compliance with the License. +@rem You may obtain a copy of the License at +@rem +@rem https://www.apache.org/licenses/LICENSE-2.0 +@rem +@rem Unless required by applicable law or agreed to in writing, software +@rem distributed under the License is distributed on an "AS IS" BASIS, +@rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +@rem See the License for the specific language governing permissions and +@rem limitations under the License. +@rem +@rem SPDX-License-Identifier: Apache-2.0 +@rem + +@if "%DEBUG%"=="" @echo off +@rem ########################################################################## +@rem +@rem gradlew startup script for Windows +@rem +@rem ########################################################################## + +@rem Set local scope for the variables, and ensure extensions are enabled +setlocal EnableExtensions + +set DIRNAME=%~dp0 +if "%DIRNAME%"=="" set DIRNAME=. +@rem This is normally unused +set APP_BASE_NAME=%~n0 +set APP_HOME=%DIRNAME% + +@rem Resolve any "." and ".." in APP_HOME to make it shorter. +for %%i in ("%APP_HOME%") do set APP_HOME=%%~fi + +@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. +set DEFAULT_JVM_OPTS="-Xmx64m" "-Xms64m" + +@rem Find java.exe +if defined JAVA_HOME goto findJavaFromJavaHome + +set JAVA_EXE=java.exe +%JAVA_EXE% -version >NUL 2>&1 +if %ERRORLEVEL% equ 0 goto execute + +echo. 1>&2 +echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. 1>&2 +echo. 1>&2 +echo Please set the JAVA_HOME variable in your environment to match the 1>&2 +echo location of your Java installation. 1>&2 + +"%COMSPEC%" /c exit 1 + +:findJavaFromJavaHome +set JAVA_HOME=%JAVA_HOME:"=% +set JAVA_EXE=%JAVA_HOME%/bin/java.exe + +if exist "%JAVA_EXE%" goto execute + +echo. 1>&2 +echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME% 1>&2 +echo. 1>&2 +echo Please set the JAVA_HOME variable in your environment to match the 1>&2 +echo location of your Java installation. 1>&2 + +"%COMSPEC%" /c exit 1 + +:execute +@rem Setup the command line + + + +@rem Execute gradlew +@rem endlocal doesn't take effect until after the line is parsed and variables are expanded +@rem which allows us to clear the local environment before executing the java command +endlocal & "%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -jar "%APP_HOME%\gradle\wrapper\gradle-wrapper.jar" %* & call :exitWithErrorLevel + +:exitWithErrorLevel +@rem Use "%COMSPEC%" /c exit to allow operators to work properly in scripts +"%COMSPEC%" /c exit %ERRORLEVEL% diff --git a/EmbeddingInversion/settings.gradle.kts b/EmbeddingInversion/settings.gradle.kts new file mode 100644 index 0000000..bad86d7 --- /dev/null +++ b/EmbeddingInversion/settings.gradle.kts @@ -0,0 +1,23 @@ +pluginManagement { + repositories { + google() + mavenCentral() + gradlePluginPortal() + } +} + +dependencyResolutionManagement { + // Repositories are declared per-module (:cli, :app) so their own blocks apply; see there. + repositories { + google() + mavenCentral() + } +} + +rootProject.name = "embedding-inversion" + +// Everything is consumed from Maven Central now: SKaiNET core + SKaiNET-transformers +// (incl. the t5 / vec2text modules) are all published at 0.36.0. No composite build. + +include(":cli") +include(":app")