kjarni-wasm 0.1.4

package/README.md

@@ -0,0 +1,123 @@
+# kjarni-wasm
+
+Local AI inference in the browser. Embeddings, classification, reranking and chat,
+running in a Web Worker so the page never freezes.
+
+```bash
+npm i kjarni-wasm
+```
+
+Or without a bundler, straight from a CDN:
+
+```html
+<script type="module">
+  import { Kjarni } from "https://cdn.jsdelivr.net/npm/kjarni-wasm@0.1.4/dist/index.js";
+</script>
+```
+
+## Quick tour
+
+```ts
+import { Kjarni } from "kjarni-wasm";
+
+const kjarni = await Kjarni.load({ encoder: "/models/minilm-l6-v2-q8.kjq" });
+
+const [doctor, physician] = await kjarni.encode(["doctor", "physician"]);
+console.log(Kjarni.cosine(doctor, physician));   // 0.8598
+
+// or in one call
+console.log(await kjarni.similarity("doctor", "physician"));
+```
+
+Every method returns a promise and the work happens on a worker thread. A model
+loads or a reply generates while the page keeps painting.
+
+## Why a worker
+
+Inference is compute-bound and synchronous. On the main thread, encoding a batch
+freezes the tab for as long as it takes, and generating a chat reply freezes it for
+seconds. That is the single most common way a browser AI feature ends up feeling
+broken.
+
+This package puts all of it on a worker and gives you promises, so you never have
+to think about it.
+
+## Loading models
+
+Models are `.kjq` files: config, tokenizer and int8 weights in one download.
+
+```ts
+const kjarni = await Kjarni.load({
+  encoder:    "/models/minilm-l6-v2-q8.kjq",
+  classifier: "/models/distilbert-sentiment-q8.kjq",
+  reranker:   "/models/ms-marco-MiniLM-L-6-v2-q8.kjq",
+  onProgress: (capability, loaded, total) => {
+    console.log(`${capability}: ${Math.round((loaded / total) * 100)}%`);
+  },
+});
+```
+
+Load only what you use. Each model is a real download, from 23 MB for MiniLM to
+several hundred for a chat model, so `onProgress` is worth wiring to a progress bar.
+
+Bytes are transferred to the worker rather than copied, which matters: copying a
+500 MB model is a visible pause and doubles peak memory.
+
+## Classification
+
+```ts
+const [result] = await kjarni.classify("This is wonderful");
+console.log(result.label, result.score);   // POSITIVE 0.9999
+
+const batch = await kjarni.classify(["Great service", "Worst purchase ever"]);
+console.log(await kjarni.labels());        // ["NEGATIVE", "POSITIVE"]
+```
+
+## Reranking
+
+Scores query and document together, which is markedly more accurate than comparing
+embeddings. Indices point back into your input, so you keep your own ids.
+
+```ts
+const ranked = await kjarni.rerank("how do I cancel?", documents, 5);
+for (const { index, score } of ranked) console.log(score, documents[index]);
+```
+
+## Chat
+
+```ts
+const kjarni = await Kjarni.load({
+  chat: "/models/qwen05b-q8.kjq",
+  chatModelId: "qwen2.5-0.5b-instruct",
+});
+
+for await (const piece of kjarni.chat("Explain RAG in one sentence.")) {
+  output.textContent += piece;
+}
+```
+
+Tokens arrive as they are decoded, not in one lump at the end. The worker is what
+makes that useful: generation occupies its thread completely, so the same call on
+the main thread would freeze the tab until the reply finished.
+
+`temperature: 0` selects greedy decoding, which is deterministic.
+
+## Cleaning up
+
+```ts
+await kjarni.close();
+```
+
+Wasm handles are not reachable by the JavaScript garbage collector, so dropping the
+reference leaks the weights. `close()` frees the models and stops the worker.
+
+## Requirements
+
+A browser with WebAssembly and module workers: Chrome and Edge 80+, Firefox 114+,
+Safari 15+. No WebGPU required, inference runs on the CPU.
+
+## Making your own `.kjq`
+
+See [the format guide](../scripts/README.md). Any BERT-family encoder,
+cross-encoder, classifier or small decoder from Hugging Face converts with one
+script.

package/dist/index.d.ts

@@ -0,0 +1,50 @@
+import type { Capability, ClassifyResult, ModelSource, RankedDocument } from "./protocol.js";
+export type { ClassifyResult, ModelSource, RankedDocument };
+export interface LoadOptions {
+    /** Embeddings and semantic similarity. A `.kjq` file. */
+    encoder?: ModelSource;
+    /** Sentiment, emotion, toxicity. A `.kjq` file. */
+    classifier?: ModelSource;
+    /** Cross-encoder reranking. A `.kjq` file. */
+    reranker?: ModelSource;
+    /** A local language model. A `.kjq` file. */
+    chat?: ModelSource;
+    /** Chat only: the registry id whose prompt template should be used. */
+    chatModelId?: string;
+    /** Called as each model is fetched, for progress UI. */
+    onProgress?: (capability: Capability, loaded: number, total: number) => void;
+}
+export interface ChatOptions {
+    maxNewTokens?: number;
+    /** 0 selects greedy decoding, which is deterministic. */
+    temperature?: number;
+}
+export declare class Kjarni {
+    #private;
+    private constructor();
+    /** Starts a worker and loads whichever models you name. */
+    static load(options: LoadOptions): Promise<Kjarni>;
+    /** One vector per input, in the order given. */
+    encode(texts: string[], normalize?: boolean): Promise<Float32Array[]>;
+    /** Cosine similarity between two texts, from -1 to 1. */
+    similarity(a: string, b: string): Promise<number>;
+    /** The best label for each input. */
+    classify(texts: string | string[]): Promise<ClassifyResult[]>;
+    /** Every label this classifier can predict, in index order. */
+    labels(): Promise<string[]>;
+    /** Documents reordered by relevance, best first. Indices point into your input. */
+    rerank(query: string, documents: string[], limit?: number): Promise<RankedDocument[]>;
+    /** The chat model's context window, in tokens. */
+    contextSize(): Promise<number>;
+    /**
+     * Generates a reply, yielding text as it arrives.
+     *
+     *   for await (const piece of kjarni.chat("Explain RAG")) out.textContent += piece;
+     */
+    chat(prompt: string, options?: ChatOptions): AsyncGenerator<string, void, unknown>;
+    /** Releases the models and stops the worker. */
+    close(): Promise<void>;
+    /** Cosine similarity between two vectors. */
+    static cosine(a: Float32Array, b: Float32Array): number;
+}
+export default Kjarni;

package/dist/index.js

@@ -0,0 +1,165 @@
+// Local AI inference in the browser, off the main thread.
+//
+//   const kjarni = await Kjarni.load({ encoder: "/models/minilm.kjq" });
+//   const [a, b] = await kjarni.encode(["doctor", "physician"]);
+//   console.log(Kjarni.cosine(a, b));   // 0.8598
+//
+// Every method returns a promise and the work happens in a worker, so the page
+// keeps painting while a model loads or a reply generates.
+async function toBytes(src, capability, onProgress) {
+    if (src instanceof Uint8Array)
+        return src;
+    if (src instanceof ArrayBuffer)
+        return new Uint8Array(src);
+    const response = await fetch(String(src));
+    if (!response.ok) {
+        throw new Error(`fetching ${src} failed: HTTP ${response.status}`);
+    }
+    // Models run to hundreds of megabytes, so report progress rather than leaving
+    // the page silent for a minute. Falls back to a plain read when the server
+    // sends no length or the body cannot be streamed.
+    const total = Number(response.headers.get("content-length") ?? 0);
+    if (!onProgress || !total || !response.body) {
+        return new Uint8Array(await response.arrayBuffer());
+    }
+    const reader = response.body.getReader();
+    const chunks = [];
+    let loaded = 0;
+    for (;;) {
+        const { done, value } = await reader.read();
+        if (done)
+            break;
+        chunks.push(value);
+        loaded += value.length;
+        onProgress(capability, loaded, total);
+    }
+    const bytes = new Uint8Array(loaded);
+    let offset = 0;
+    for (const chunk of chunks) {
+        bytes.set(chunk, offset);
+        offset += chunk.length;
+    }
+    return bytes;
+}
+export class Kjarni {
+    #worker;
+    #next = 1;
+    #pending = new Map();
+    constructor(worker) {
+        this.#worker = worker;
+        this.#worker.onmessage = (event) => {
+            const msg = event.data;
+            const entry = this.#pending.get(msg.id);
+            if (!entry)
+                return;
+            if (msg.kind === "token") {
+                entry.onToken?.(msg.text);
+                return; // more to come; the promise settles on ok or err
+            }
+            this.#pending.delete(msg.id);
+            if (msg.kind === "err")
+                entry.reject(new Error(msg.message));
+            else
+                entry.resolve(msg.value);
+        };
+    }
+    /** Starts a worker and loads whichever models you name. */
+    static async load(options) {
+        const worker = new Worker(new URL("./worker.js", import.meta.url), { type: "module" });
+        const kjarni = new Kjarni(worker);
+        const capabilities = ["encoder", "classifier", "reranker", "chat"];
+        for (const capability of capabilities) {
+            const source = options[capability];
+            if (!source)
+                continue;
+            const bytes = await toBytes(source, capability, options.onProgress);
+            await kjarni.#send({ kind: "load", capability, bytes, modelId: options.chatModelId }, 
+            // The buffer is handed to the worker rather than copied. A 500MB model
+            // copied per load is a real pause and a doubled peak memory footprint.
+            [bytes.buffer]);
+        }
+        return kjarni;
+    }
+    #send(req, transfer = [], onToken) {
+        const id = this.#next++;
+        return new Promise((resolve, reject) => {
+            this.#pending.set(id, { resolve, reject, onToken });
+            this.#worker.postMessage({ ...req, id }, transfer);
+        });
+    }
+    /** One vector per input, in the order given. */
+    async encode(texts, normalize = true) {
+        return (await this.#send({ kind: "encode", texts, normalize }));
+    }
+    /** Cosine similarity between two texts, from -1 to 1. */
+    async similarity(a, b) {
+        const [va, vb] = await this.encode([a, b]);
+        return Kjarni.cosine(va, vb);
+    }
+    /** The best label for each input. */
+    async classify(texts) {
+        const batch = Array.isArray(texts) ? texts : [texts];
+        return (await this.#send({ kind: "classify", texts: batch }));
+    }
+    /** Every label this classifier can predict, in index order. */
+    async labels() {
+        return (await this.#send({ kind: "labels" }));
+    }
+    /** Documents reordered by relevance, best first. Indices point into your input. */
+    async rerank(query, documents, limit = documents.length) {
+        return (await this.#send({ kind: "rerank", query, documents, limit }));
+    }
+    /** The chat model's context window, in tokens. */
+    async contextSize() {
+        return (await this.#send({ kind: "contextSize" }));
+    }
+    /**
+     * Generates a reply, yielding text as it arrives.
+     *
+     *   for await (const piece of kjarni.chat("Explain RAG")) out.textContent += piece;
+     */
+    async *chat(prompt, options = {}) {
+        const queue = [];
+        let wake = null;
+        const done = this.#send({
+            kind: "chat",
+            prompt,
+            maxNewTokens: options.maxNewTokens ?? 256,
+            temperature: options.temperature ?? 0,
+        }, [], (text) => {
+            queue.push(text);
+            wake?.();
+        });
+        let finished = false;
+        const settled = done.then(() => { finished = true; wake?.(); }, (e) => { finished = true; wake?.(); throw e; });
+        while (!finished || queue.length) {
+            if (queue.length) {
+                yield queue.shift();
+                continue;
+            }
+            await new Promise((r) => (wake = r));
+        }
+        await settled; // surfaces a generation error rather than ending quietly
+    }
+    /** Releases the models and stops the worker. */
+    async close() {
+        try {
+            await this.#send({ kind: "close" });
+        }
+        finally {
+            this.#worker.terminate();
+        }
+    }
+    /** Cosine similarity between two vectors. */
+    static cosine(a, b) {
+        let dot = 0, na = 0, nb = 0;
+        for (let i = 0; i < a.length; i++) {
+            dot += a[i] * b[i];
+            na += a[i] * a[i];
+            nb += b[i] * b[i];
+        }
+        return dot / (Math.sqrt(na) * Math.sqrt(nb) || 1e-9);
+    }
+}
+export default Kjarni;
+//# sourceMappingURL=index.js.map

package/dist/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,0DAA0D;AAC1D,EAAE;AACF,yEAAyE;AACzE,iEAAiE;AACjE,kDAAkD;AAClD,EAAE;AACF,+EAA+E;AAC/E,2DAA2D;AAkC3D,KAAK,UAAU,OAAO,CACpB,GAAgB,EAChB,UAAsB,EACtB,UAAsC;IAEtC,IAAI,GAAG,YAAY,UAAU;QAAE,OAAO,GAAG,CAAC;IAC1C,IAAI,GAAG,YAAY,WAAW;QAAE,OAAO,IAAI,UAAU,CAAC,GAAG,CAAC,CAAC;IAE3D,MAAM,QAAQ,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,GAAG,CAAC,CAAC,CAAC;IAC1C,IAAI,CAAC,QAAQ,CAAC,EAAE,EAAE,CAAC;QACjB,MAAM,IAAI,KAAK,CAAC,YAAY,GAAG,iBAAiB,QAAQ,CAAC,MAAM,EAAE,CAAC,CAAC;IACrE,CAAC;IAED,8EAA8E;IAC9E,2EAA2E;IAC3E,kDAAkD;IAClD,MAAM,KAAK,GAAG,MAAM,CAAC,QAAQ,CAAC,OAAO,CAAC,GAAG,CAAC,gBAAgB,CAAC,IAAI,CAAC,CAAC,CAAC;IAClE,IAAI,CAAC,UAAU,IAAI,CAAC,KAAK,IAAI,CAAC,QAAQ,CAAC,IAAI,EAAE,CAAC;QAC5C,OAAO,IAAI,UAAU,CAAC,MAAM,QAAQ,CAAC,WAAW,EAAE,CAAC,CAAC;IACtD,CAAC;IAED,MAAM,MAAM,GAAG,QAAQ,CAAC,IAAI,CAAC,SAAS,EAAE,CAAC;IACzC,MAAM,MAAM,GAAiB,EAAE,CAAC;IAChC,IAAI,MAAM,GAAG,CAAC,CAAC;IACf,SAAS,CAAC;QACR,MAAM,EAAE,IAAI,EAAE,KAAK,EAAE,GAAG,MAAM,MAAM,CAAC,IAAI,EAAE,CAAC;QAC5C,IAAI,IAAI;YAAE,MAAM;QAChB,MAAM,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;QACnB,MAAM,IAAI,KAAK,CAAC,MAAM,CAAC;QACvB,UAAU,CAAC,UAAU,EAAE,MAAM,EAAE,KAAK,CAAC,CAAC;IACxC,CAAC;IAED,MAAM,KAAK,GAAG,IAAI,UAAU,CAAC,MAAM,CAAC,CAAC;IACrC,IAAI,MAAM,GAAG,CAAC,CAAC;IACf,KAAK,MAAM,KAAK,IAAI,MAAM,EAAE,CAAC;QAC3B,KAAK,CAAC,GAAG,CAAC,KAAK,EAAE,MAAM,CAAC,CAAC;QACzB,MAAM,IAAI,KAAK,CAAC,MAAM,CAAC;IACzB,CAAC;IACD,OAAO,KAAK,CAAC;AACf,CAAC;AAaD,MAAM,OAAO,MAAM;IACjB,OAAO,CAAS;IAChB,KAAK,GAAG,CAAC,CAAC;IACV,QAAQ,GAAG,IAAI,GAAG,EAGf,CAAC;IAEJ,YAAoB,MAAc;QAChC,IAAI,CAAC,OAAO,GAAG,MAAM,CAAC;QACtB,IAAI,CAAC,OAAO,CAAC,SAAS,GAAG,CAAC,KAA6B,EAAE,EAAE;YACzD,MAAM,GAAG,GAAG,KAAK,CAAC,IAAI,CAAC;YACvB,MAAM,KAAK,GAAG,IAAI,CAAC,QAAQ,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;YACxC,IAAI,CAAC,KAAK;gBAAE,OAAO;YAEnB,IAAI,GAAG,CAAC,IAAI,KAAK,OAAO,EAAE,CAAC;gBACzB,KAAK,CAAC,OAAO,EAAE,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;gBAC1B,OAAO,CAAC,iDAAiD;YAC3D,CAAC;YACD,IAAI,CAAC,QAAQ,CAAC,MAAM,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;YAC7B,IAAI,GAAG,CAAC,IAAI,KAAK,KAAK;gBAAE,KAAK,CAAC,MAAM,CAAC,IAAI,KAAK,CAAC,GAAG,CAAC,OAAO,CAAC,CAAC,CAAC;;gBACxD,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC;QAChC,CAAC,CAAC;IACJ,CAAC;IAED,2DAA2D;IAC3D,MAAM,CAAC,KAAK,CAAC,IAAI,CAAC,OAAoB;QACpC,MAAM,MAAM,GAAG,IAAI,MAAM,CAAC,IAAI,GAAG,CAAC,aAAa,EAAE,MAAM,CAAC,IAAI,CAAC,GAAG,CAAC,EAAE,EAAE,IAAI,EAAE,QAAQ,EAAE,CAAC,CAAC;QACvF,MAAM,MAAM,GAAG,IAAI,MAAM,CAAC,MAAM,CAAC,CAAC;QAElC,MAAM,YAAY,GAAiB,CAAC,SAAS,EAAE,YAAY,EAAE,UAAU,EAAE,MAAM,CAAC,CAAC;QACjF,KAAK,MAAM,UAAU,IAAI,YAAY,EAAE,CAAC;YACtC,MAAM,MAAM,GAAG,OAAO,CAAC,UAAU,CAAC,CAAC;YACnC,IAAI,CAAC,MAAM;gBAAE,SAAS;YACtB,MAAM,KAAK,GAAG,MAAM,OAAO,CAAC,MAAM,EAAE,UAAU,EAAE,OAAO,CAAC,UAAU,CAAC,CAAC;YACpE,MAAM,MAAM,CAAC,KAAK,CAChB,EAAE,IAAI,EAAE,MAAM,EAAE,UAAU,EAAE,KAAK,EAAE,OAAO,EAAE,OAAO,CAAC,WAAW,EAAE;YACjE,uEAAuE;YACvE,uEAAuE;YACvE,CAAC,KAAK,CAAC,MAAM,CAAC,CACf,CAAC;QACJ,CAAC;QACD,OAAO,MAAM,CAAC;IAChB,CAAC;IAED,KAAK,CAAC,GAAmB,EAAE,WAA2B,EAAE,EAAE,OAA6B;QACrF,MAAM,EAAE,GAAG,IAAI,CAAC,KAAK,EAAE,CAAC;QACxB,OAAO,IAAI,OAAO,CAAU,CAAC,OAAO,EAAE,MAAM,EAAE,EAAE;YAC9C,IAAI,CAAC,QAAQ,CAAC,GAAG,CAAC,EAAE,EAAE,EAAE,OAAO,EAAE,MAAM,EAAE,OAAO,EAAE,CAAC,CAAC;YACpD,IAAI,CAAC,OAAO,CAAC,WAAW,CAAC,EAAE,GAAG,GAAG,EAAE,EAAE,EAAa,EAAE,QAAQ,CAAC,CAAC;QAChE,CAAC,CAAC,CAAC;IACL,CAAC;IAED,gDAAgD;IAChD,KAAK,CAAC,MAAM,CAAC,KAAe,EAAE,SAAS,GAAG,IAAI;QAC5C,OAAO,CAAC,MAAM,IAAI,CAAC,KAAK,CAAC,EAAE,IAAI,EAAE,QAAQ,EAAE,KAAK,EAAE,SAAS,EAAE,CAAC,CAAmB,CAAC;IACpF,CAAC;IAED,yDAAyD;IACzD,KAAK,CAAC,UAAU,CAAC,CAAS,EAAE,CAAS;QACnC,MAAM,CAAC,EAAE,EAAE,EAAE,CAAC,GAAG,MAAM,IAAI,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAC3C,OAAO,MAAM,CAAC,MAAM,CAAC,EAAE,EAAE,EAAE,CAAC,CAAC;IAC/B,CAAC;IAED,qCAAqC;IACrC,KAAK,CAAC,QAAQ,CAAC,KAAwB;QACrC,MAAM,KAAK,GAAG,KAAK,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC;QACrD,OAAO,CAAC,MAAM,IAAI,CAAC,KAAK,CAAC,EAAE,IAAI,EAAE,UAAU,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAqB,CAAC;IACpF,CAAC;IAED,+DAA+D;IAC/D,KAAK,CAAC,MAAM;QACV,OAAO,CAAC,MAAM,IAAI,CAAC,KAAK,CAAC,EAAE,IAAI,EAAE,QAAQ,EAAE,CAAC,CAAa,CAAC;IAC5D,CAAC;IAED,mFAAmF;IACnF,KAAK,CAAC,MAAM,CAAC,KAAa,EAAE,SAAmB,EAAE,KAAK,GAAG,SAAS,CAAC,MAAM;QACvE,OAAO,CAAC,MAAM,IAAI,CAAC,KAAK,CAAC,EAAE,IAAI,EAAE,QAAQ,EAAE,KAAK,EAAE,SAAS,EAAE,KAAK,EAAE,CAAC,CAAqB,CAAC;IAC7F,CAAC;IAED,kDAAkD;IAClD,KAAK,CAAC,WAAW;QACf,OAAO,CAAC,MAAM,IAAI,CAAC,KAAK,CAAC,EAAE,IAAI,EAAE,aAAa,EAAE,CAAC,CAAW,CAAC;IAC/D,CAAC;IAED;;;;OAIG;IACH,KAAK,CAAC,CAAC,IAAI,CAAC,MAAc,EAAE,UAAuB,EAAE;QACnD,MAAM,KAAK,GAAa,EAAE,CAAC;QAC3B,IAAI,IAAI,GAAwB,IAAI,CAAC;QAErC,MAAM,IAAI,GAAG,IAAI,CAAC,KAAK,CACrB;YACE,IAAI,EAAE,MAAM;YACZ,MAAM;YACN,YAAY,EAAE,OAAO,CAAC,YAAY,IAAI,GAAG;YACzC,WAAW,EAAE,OAAO,CAAC,WAAW,IAAI,CAAC;SACtC,EACD,EAAE,EACF,CAAC,IAAI,EAAE,EAAE;YACP,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;YACjB,IAAI,EAAE,EAAE,CAAC;QACX,CAAC,CACF,CAAC;QAEF,IAAI,QAAQ,GAAG,KAAK,CAAC;QACrB,MAAM,OAAO,GAAG,IAAI,CAAC,IAAI,CACvB,GAAG,EAAE,GAAG,QAAQ,GAAG,IAAI,CAAC,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,CAAC,EACpC,CAAC,CAAC,EAAE,EAAE,GAAG,QAAQ,GAAG,IAAI,CAAC,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAC/C,CAAC;QAEF,OAAO,CAAC,QAAQ,IAAI,KAAK,CAAC,MAAM,EAAE,CAAC;YACjC,IAAI,KAAK,CAAC,MAAM,EAAE,CAAC;gBACjB,MAAM,KAAK,CAAC,KAAK,EAAG,CAAC;gBACrB,SAAS;YACX,CAAC;YACD,MAAM,IAAI,OAAO,CAAO,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,IAAI,GAAG,CAAC,CAAC,CAAC,CAAC;QAC7C,CAAC;QACD,MAAM,OAAO,CAAC,CAAG,yDAAyD;IAC5E,CAAC;IAED,gDAAgD;IAChD,KAAK,CAAC,KAAK;QACT,IAAI,CAAC;YACH,MAAM,IAAI,CAAC,KAAK,CAAC,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAC;QACtC,CAAC;gBAAS,CAAC;YACT,IAAI,CAAC,OAAO,CAAC,SAAS,EAAE,CAAC;QAC3B,CAAC;IACH,CAAC;IAED,6CAA6C;IAC7C,MAAM,CAAC,MAAM,CAAC,CAAe,EAAE,CAAe;QAC5C,IAAI,GAAG,GAAG,CAAC,EAAE,EAAE,GAAG,CAAC,EAAE,EAAE,GAAG,CAAC,CAAC;QAC5B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YAClC,GAAG,IAAI,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC;YACnB,EAAE,IAAI,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC;YAClB,EAAE,IAAI,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC;QACpB,CAAC;QACD,OAAO,GAAG,GAAG,CAAC,IAAI,CAAC,IAAI,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,IAAI,CAAC,EAAE,CAAC,IAAI,IAAI,CAAC,CAAC;IACvD,CAAC;CACF;AAED,eAAe,MAAM,CAAC"}

package/dist/protocol.d.ts

@@ -0,0 +1,63 @@
+/** A model to load, either already in memory or somewhere to fetch it from. */
+export type ModelSource = string | URL | ArrayBuffer | Uint8Array;
+export type Capability = "encoder" | "classifier" | "reranker" | "chat";
+export interface ClassifyResult {
+    label: string;
+    score: number;
+    index: number;
+}
+export interface RankedDocument {
+    index: number;
+    score: number;
+    text: string;
+}
+export type Request = {
+    id: number;
+    kind: "load";
+    capability: Capability;
+    bytes: Uint8Array;
+    modelId?: string;
+} | {
+    id: number;
+    kind: "encode";
+    texts: string[];
+    normalize: boolean;
+} | {
+    id: number;
+    kind: "classify";
+    texts: string[];
+} | {
+    id: number;
+    kind: "labels";
+} | {
+    id: number;
+    kind: "rerank";
+    query: string;
+    documents: string[];
+    limit: number;
+} | {
+    id: number;
+    kind: "chat";
+    prompt: string;
+    maxNewTokens: number;
+    temperature: number;
+} | {
+    id: number;
+    kind: "contextSize";
+} | {
+    id: number;
+    kind: "close";
+};
+export type Response = {
+    id: number;
+    kind: "ok";
+    value: unknown;
+} | {
+    id: number;
+    kind: "token";
+    text: string;
+} | {
+    id: number;
+    kind: "err";
+    message: string;
+};

package/dist/protocol.js

@@ -0,0 +1,8 @@
+// Messages between the page and the worker.
+//
+// Every request carries an id and is answered by exactly one Ok or Err with the
+// same id, except `chat`, which streams any number of Token messages before its
+// Ok. Keeping that contract in one file is what lets both sides be type-checked
+// against the same shapes rather than agreeing by convention.
+export {};
+//# sourceMappingURL=protocol.js.map

package/dist/protocol.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"protocol.js","sourceRoot":"","sources":["../src/protocol.ts"],"names":[],"mappings":"AAAA,4CAA4C;AAC5C,EAAE;AACF,gFAAgF;AAChF,gFAAgF;AAChF,gFAAgF;AAChF,8DAA8D"}

package/dist/worker.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/worker.js

@@ -0,0 +1,109 @@
+// The worker. Everything expensive happens here so the page never janks.
+//
+// Inference in this build is synchronous: `encode` and `generate` occupy the
+// thread until they return. On the main thread that is a frozen tab for the
+// duration, which for a chat reply is seconds. Here it costs nothing, because
+// this thread has nothing else to do.
+import init, { WasmModel, WasmClassifier, WasmReranker, WasmChat, } from "../pkg/kjarni_wasm.js";
+let model = null;
+let classifier = null;
+let reranker = null;
+let chat = null;
+let ready = null;
+const post = (msg) => self.postMessage(msg);
+/** Loads the wasm binary once, however many times this is called. */
+function ensureWasm() {
+    ready ??= init();
+    return ready;
+}
+function handle(req) {
+    switch (req.kind) {
+        case "load": {
+            // Loading a second model of the same kind replaces the first, and the old
+            // one has to be freed explicitly: wasm-bindgen handles cannot be collected
+            // by the JS garbage collector, so dropping the reference leaks the weights.
+            switch (req.capability) {
+                case "encoder":
+                    model?.free();
+                    model = WasmModel.from_quantized(req.bytes);
+                    return null;
+                case "classifier":
+                    classifier?.free();
+                    classifier = WasmClassifier.load(req.bytes);
+                    return null;
+                case "reranker":
+                    reranker?.free();
+                    reranker = WasmReranker.load(req.bytes);
+                    return null;
+                case "chat":
+                    chat?.free();
+                    chat = WasmChat.load(req.bytes, req.modelId ?? null);
+                    return null;
+            }
+        }
+        case "encode": {
+            if (!model)
+                throw new Error("no encoder loaded; call load({ encoder }) first");
+            // The binding returns every vector concatenated; split it back out here so
+            // callers get one array per input rather than doing the arithmetic.
+            const flat = model.encode(req.texts, req.normalize);
+            const dim = flat.length / req.texts.length;
+            const out = [];
+            for (let i = 0; i < req.texts.length; i++) {
+                out.push(flat.slice(i * dim, (i + 1) * dim));
+            }
+            return out;
+        }
+        case "classify": {
+            if (!classifier)
+                throw new Error("no classifier loaded; call load({ classifier }) first");
+            return req.texts.length === 1
+                ? [classifier.classify(req.texts[0])]
+                : classifier.classify_batch(req.texts);
+        }
+        case "labels":
+            if (!classifier)
+                throw new Error("no classifier loaded");
+            return classifier.labels();
+        case "rerank":
+            if (!reranker)
+                throw new Error("no reranker loaded; call load({ reranker }) first");
+            return reranker.rerank(req.query, req.documents, req.limit);
+        case "contextSize":
+            if (!chat)
+                throw new Error("no chat model loaded");
+            return chat.context_size();
+        case "chat": {
+            if (!chat)
+                throw new Error("no chat model loaded; call load({ chat }) first");
+            // Genuinely incremental: the callback fires between decoding steps, and
+            // each postMessage reaches the page while the model is still generating.
+            // This is the reason the whole wrapper is worth having, since the same
+            // call on the main thread would freeze the tab until the reply finished.
+            chat.generate_stream(req.prompt, req.maxNewTokens, req.temperature, (piece) => {
+                post({ id: req.id, kind: "token", text: piece });
+            });
+            return null;
+        }
+        case "close": {
+            model?.free();
+            classifier?.free();
+            reranker?.free();
+            chat?.free();
+            model = classifier = reranker = chat = null;
+            return null;
+        }
+    }
+}
+self.onmessage = async (event) => {
+    const req = event.data;
+    try {
+        await ensureWasm();
+        const value = handle(req);
+        post({ id: req.id, kind: "ok", value });
+    }
+    catch (e) {
+        post({ id: req.id, kind: "err", message: e instanceof Error ? e.message : String(e) });
+    }
+};
+//# sourceMappingURL=worker.js.map

package/dist/worker.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"worker.js","sourceRoot":"","sources":["../src/worker.ts"],"names":[],"mappings":"AAAA,yEAAyE;AACzE,EAAE;AACF,6EAA6E;AAC7E,4EAA4E;AAC5E,8EAA8E;AAC9E,sCAAsC;AAEtC,OAAO,IAAI,EAAE,EACX,SAAS,EACT,cAAc,EACd,YAAY,EACZ,QAAQ,GACT,MAAM,uBAAuB,CAAC;AAI/B,IAAI,KAAK,GAAqB,IAAI,CAAC;AACnC,IAAI,UAAU,GAA0B,IAAI,CAAC;AAC7C,IAAI,QAAQ,GAAwB,IAAI,CAAC;AACzC,IAAI,IAAI,GAAoB,IAAI,CAAC;AACjC,IAAI,KAAK,GAA4B,IAAI,CAAC;AAE1C,MAAM,IAAI,GAAG,CAAC,GAAa,EAAE,EAAE,CAAE,IAA0B,CAAC,WAAW,CAAC,GAAG,CAAC,CAAC;AAE7E,qEAAqE;AACrE,SAAS,UAAU;IACjB,KAAK,KAAK,IAAI,EAAE,CAAC;IACjB,OAAO,KAAK,CAAC;AACf,CAAC;AAED,SAAS,MAAM,CAAC,GAAY;IAC1B,QAAQ,GAAG,CAAC,IAAI,EAAE,CAAC;QACjB,KAAK,MAAM,CAAC,CAAC,CAAC;YACZ,0EAA0E;YAC1E,2EAA2E;YAC3E,4EAA4E;YAC5E,QAAQ,GAAG,CAAC,UAAU,EAAE,CAAC;gBACvB,KAAK,SAAS;oBACZ,KAAK,EAAE,IAAI,EAAE,CAAC;oBACd,KAAK,GAAG,SAAS,CAAC,cAAc,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC;oBAC5C,OAAO,IAAI,CAAC;gBACd,KAAK,YAAY;oBACf,UAAU,EAAE,IAAI,EAAE,CAAC;oBACnB,UAAU,GAAG,cAAc,CAAC,IAAI,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC;oBAC5C,OAAO,IAAI,CAAC;gBACd,KAAK,UAAU;oBACb,QAAQ,EAAE,IAAI,EAAE,CAAC;oBACjB,QAAQ,GAAG,YAAY,CAAC,IAAI,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC;oBACxC,OAAO,IAAI,CAAC;gBACd,KAAK,MAAM;oBACT,IAAI,EAAE,IAAI,EAAE,CAAC;oBACb,IAAI,GAAG,QAAQ,CAAC,IAAI,CAAC,GAAG,CAAC,KAAK,EAAE,GAAG,CAAC,OAAO,IAAI,IAAI,CAAC,CAAC;oBACrD,OAAO,IAAI,CAAC;YAChB,CAAC;QACH,CAAC;QAED,KAAK,QAAQ,CAAC,CAAC,CAAC;YACd,IAAI,CAAC,KAAK;gBAAE,MAAM,IAAI,KAAK,CAAC,iDAAiD,CAAC,CAAC;YAC/E,2EAA2E;YAC3E,oEAAoE;YACpE,MAAM,IAAI,GAAG,KAAK,CAAC,MAAM,CAAC,GAAG,CAAC,KAAK,EAAE,GAAG,CAAC,SAAS,CAAC,CAAC;YACpD,MAAM,GAAG,GAAG,IAAI,CAAC,MAAM,GAAG,GAAG,CAAC,KAAK,CAAC,MAAM,CAAC;YAC3C,MAAM,GAAG,GAAmB,EAAE,CAAC;YAC/B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,GAAG,CAAC,KAAK,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;gBAC1C,GAAG,CAAC,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC,GAAG,GAAG,EAAE,CAAC,CAAC,GAAG,CAAC,CAAC,GAAG,GAAG,CAAC,CAAC,CAAC;YAC/C,CAAC;YACD,OAAO,GAAG,CAAC;QACb,CAAC;QAED,KAAK,UAAU,CAAC,CAAC,CAAC;YAChB,IAAI,CAAC,UAAU;gBAAE,MAAM,IAAI,KAAK,CAAC,uDAAuD,CAAC,CAAC;YAC1F,OAAO,GAAG,CAAC,KAAK,CAAC,MAAM,KAAK,CAAC;gBAC3B,CAAC,CAAC,CAAC,UAAU,CAAC,QAAQ,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAC;gBACrC,CAAC,CAAC,UAAU,CAAC,cAAc,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC;QAC3C,CAAC;QAED,KAAK,QAAQ;YACX,IAAI,CAAC,UAAU;gBAAE,MAAM,IAAI,KAAK,CAAC,sBAAsB,CAAC,CAAC;YACzD,OAAO,UAAU,CAAC,MAAM,EAAE,CAAC;QAE7B,KAAK,QAAQ;YACX,IAAI,CAAC,QAAQ;gBAAE,MAAM,IAAI,KAAK,CAAC,mDAAmD,CAAC,CAAC;YACpF,OAAO,QAAQ,CAAC,MAAM,CAAC,GAAG,CAAC,KAAK,EAAE,GAAG,CAAC,SAAS,EAAE,GAAG,CAAC,KAAK,CAAC,CAAC;QAE9D,KAAK,aAAa;YAChB,IAAI,CAAC,IAAI;gBAAE,MAAM,IAAI,KAAK,CAAC,sBAAsB,CAAC,CAAC;YACnD,OAAO,IAAI,CAAC,YAAY,EAAE,CAAC;QAE7B,KAAK,MAAM,CAAC,CAAC,CAAC;YACZ,IAAI,CAAC,IAAI;gBAAE,MAAM,IAAI,KAAK,CAAC,iDAAiD,CAAC,CAAC;YAC9E,wEAAwE;YACxE,yEAAyE;YACzE,uEAAuE;YACvE,yEAAyE;YACzE,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,MAAM,EAAE,GAAG,CAAC,YAAY,EAAE,GAAG,CAAC,WAAW,EAAE,CAAC,KAAa,EAAE,EAAE;gBACpF,IAAI,CAAC,EAAE,EAAE,EAAE,GAAG,CAAC,EAAE,EAAE,IAAI,EAAE,OAAO,EAAE,IAAI,EAAE,KAAK,EAAE,CAAC,CAAC;YACnD,CAAC,CAAC,CAAC;YACH,OAAO,IAAI,CAAC;QACd,CAAC;QAED,KAAK,OAAO,CAAC,CAAC,CAAC;YACb,KAAK,EAAE,IAAI,EAAE,CAAC;YACd,UAAU,EAAE,IAAI,EAAE,CAAC;YACnB,QAAQ,EAAE,IAAI,EAAE,CAAC;YACjB,IAAI,EAAE,IAAI,EAAE,CAAC;YACb,KAAK,GAAG,UAAU,GAAG,QAAQ,GAAG,IAAI,GAAG,IAAI,CAAC;YAC5C,OAAO,IAAI,CAAC;QACd,CAAC;IACH,CAAC;AACH,CAAC;AAED,IAAI,CAAC,SAAS,GAAG,KAAK,EAAE,KAA4B,EAAE,EAAE;IACtD,MAAM,GAAG,GAAG,KAAK,CAAC,IAAI,CAAC;IACvB,IAAI,CAAC;QACH,MAAM,UAAU,EAAE,CAAC;QACnB,MAAM,KAAK,GAAG,MAAM,CAAC,GAAG,CAAC,CAAC;QAC1B,IAAI,CAAC,EAAE,EAAE,EAAE,GAAG,CAAC,EAAE,EAAE,IAAI,EAAE,IAAI,EAAE,KAAK,EAAE,CAAC,CAAC;IAC1C,CAAC;IAAC,OAAO,CAAC,EAAE,CAAC;QACX,IAAI,CAAC,EAAE,EAAE,EAAE,GAAG,CAAC,EAAE,EAAE,IAAI,EAAE,KAAK,EAAE,OAAO,EAAE,CAAC,YAAY,KAAK,CAAC,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC;IACzF,CAAC;AACH,CAAC,CAAC"}

package/package.json

@@ -0,0 +1,46 @@
+{
+  "name": "kjarni-wasm",
+  "version": "0.1.4",
+  "description": "Local AI inference in the browser: embeddings, classification, reranking, semantic search and chat, running in a Web Worker.",
+  "type": "module",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "default": "./dist/index.js"
+    }
+  },
+  "files": [
+    "dist",
+    "pkg"
+  ],
+  "sideEffects": false,
+  "scripts": {
+    "prebuild": "rm -rf pkg && cp -r ../pkg pkg && rm -f pkg/.gitignore",
+    "build": "tsc -p tsconfig.json",
+    "prepack": "npm run build"
+  },
+  "devDependencies": {
+    "typescript": "^5.6.0"
+  },
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/olafurjohannsson/kjarni.git",
+    "directory": "crates/kjarni-wasm"
+  },
+  "homepage": "https://kjarni.ai",
+  "keywords": [
+    "ai",
+    "embeddings",
+    "webassembly",
+    "wasm",
+    "local",
+    "offline",
+    "bert",
+    "rag",
+    "semantic-search",
+    "web-worker"
+  ]
+}