@desert-ant-labs/emo 0.5.0 → 0.6.0

package/README.md

@@ -17,7 +17,7 @@ const toned = (await suggestions("go for a run", 1, { skinTone: "medium" }))[0]?
 ## Features
 
 - Pure-JS inference (no ONNX/WASM runtime); prediction is sub-millisecond
-- Suggests from a data-driven vocabulary of ~500 task/calendar/message emojis
+- Suggests from a curated vocabulary of ~800 everyday emojis (task, message, and concrete nouns)
 - Supports 23 languages (incl. CJK, Arabic, Thai, Hindi, …)
 - Model (~5.0 MB, 4-bit palettized) is fetched from the Hugging Face Hub at a **pinned revision**, then
   cached, to the **filesystem** on Node and to **Cache Storage** in the browser, so

package/dist/hub.d.ts

@@ -3,7 +3,7 @@ export declare const DEFAULT_HOST = "https://huggingface.co";
 export declare const DEFAULT_REPO = "desert-ant-labs/emo";
 /** Pinned revision of the model repo. A tag (not a bare commit SHA) so it
  * survives history rewrites/squashes on the model repo. */
-export declare const DEFAULT_REVISION = "v0.5.0";
+export declare const DEFAULT_REVISION = "v0.6.0";
 /** Resolution + caching configuration (mutate the exported `env` to change defaults). */
 export interface EmoEnv {
     /** Hugging Face host serving the model repo. */

package/dist/hub.js

@@ -3,7 +3,7 @@ export const DEFAULT_HOST = "https://huggingface.co";
 export const DEFAULT_REPO = "desert-ant-labs/emo";
 /** Pinned revision of the model repo. A tag (not a bare commit SHA) so it
  * survives history rewrites/squashes on the model repo. */
-export const DEFAULT_REVISION = "v0.5.0";
+export const DEFAULT_REVISION = "v0.6.0";
 const FILES = ["emo.safetensors", "emo_tokenizer.bin", "emo_meta.json"];
 function resolveUrl(env, name) {
     return `${env.host}/${env.repo}/resolve/${env.revision}/${name}`;

package/dist/model.d.ts

@@ -14,6 +14,10 @@ export interface EmoMeta {
     n_importance: number;
     sem_dim: number;
     sem_pad_index: number;
+    /** "transformer" runs the semantic sequence through an encoder; absent = mean-pool. */
+    arch?: string;
+    n_layers?: number;
+    n_heads?: number;
 }
 /** Loads the bundled model and runs emoji inference. Construct once and reuse. */
 export declare class EmoModel {

package/dist/model.js

@@ -31,6 +31,88 @@ function erf(x) {
     return x >= 0 ? y : -y;
 }
 const gelu = (x) => 0.5 * x * (1 + erf(x / Math.SQRT2));
+// ---- transformer helpers (semantic sequence path) -------------------------
+// S is a row-major T x inDim buffer; returns T x outDim = S @ W^T + b (W is outDim x inDim).
+function linear(S, T, inDim, W, outDim, b) {
+    const out = new Float32Array(T * outDim);
+    for (let t = 0; t < T; t++) {
+        const sb = t * inDim, ob = t * outDim;
+        for (let o = 0; o < outDim; o++) {
+            let acc = b ? b.get(o) : 0;
+            const wb = o * inDim;
+            for (let i = 0; i < inDim; i++)
+                acc += W.get(wb + i) * S[sb + i];
+            out[ob + o] = acc;
+        }
+    }
+    return out;
+}
+// in-place post-norm LayerNorm over the last dim.
+function layerNorm(S, T, d, w, b) {
+    for (let t = 0; t < T; t++) {
+        const base = t * d;
+        let mean = 0;
+        for (let c = 0; c < d; c++)
+            mean += S[base + c];
+        mean /= d;
+        let v = 0;
+        for (let c = 0; c < d; c++) {
+            const z = S[base + c] - mean;
+            v += z * z;
+        }
+        const inv = 1 / Math.sqrt(v / d + 1e-5);
+        for (let c = 0; c < d; c++)
+            S[base + c] = (S[base + c] - mean) * inv * w.get(c) + b.get(c);
+    }
+}
+// Post-norm encoder layer (gelu FFN), matching nn.TransformerEncoderLayer.
+function encoderLayer(S, T, d, H, L) {
+    const dh = d / H, scale = 1 / Math.sqrt(dh);
+    const qkv = linear(S, T, d, L.qkv_w, 3 * d, L.qkv_b); // T x 3d
+    const attn = new Float32Array(T * d);
+    const scores = new Float32Array(T);
+    for (let h = 0; h < H; h++) {
+        const off = h * dh;
+        for (let ti = 0; ti < T; ti++) {
+            const qb = ti * 3 * d + off;
+            let mx = -Infinity;
+            for (let tj = 0; tj < T; tj++) {
+                const kb = tj * 3 * d + d + off;
+                let s = 0;
+                for (let c = 0; c < dh; c++)
+                    s += qkv[qb + c] * qkv[kb + c];
+                s *= scale;
+                scores[tj] = s;
+                if (s > mx)
+                    mx = s;
+            }
+            let sum = 0;
+            for (let tj = 0; tj < T; tj++) {
+                const e = Math.exp(scores[tj] - mx);
+                scores[tj] = e;
+                sum += e;
+            }
+            const ob = ti * d + off;
+            for (let tj = 0; tj < T; tj++) {
+                const a = scores[tj] / sum, vb = tj * 3 * d + 2 * d + off;
+                for (let c = 0; c < dh; c++)
+                    attn[ob + c] += a * qkv[vb + c];
+            }
+        }
+    }
+    const o = linear(attn, T, d, L.o_w, d, L.o_b); // T x d
+    for (let i = 0; i < T * d; i++)
+        o[i] += S[i];
+    layerNorm(o, T, d, L.n1_w, L.n1_b);
+    const ff1 = linear(o, T, d, L.f1_w, L.f1_b.rows, L.f1_b); // T x ff
+    for (let i = 0; i < ff1.length; i++)
+        ff1[i] = gelu(ff1[i]);
+    const ff2 = linear(ff1, T, L.f1_b.rows, L.f2_w, d, L.f2_b); // T x d
+    for (let i = 0; i < T * d; i++)
+        ff2[i] += o[i];
+    layerNorm(ff2, T, d, L.n2_w, L.n2_b);
+    return ff2;
+}
 // Minimal safetensors reader: u64 header length, JSON header, then raw tensor bytes.
 // Each weight is either raw F32, or a 4-bit k-means palette stored as a packed U8 index
 // tensor plus a "<name>.palette" F32 tensor (logical 2-D shape kept in __metadata__).
@@ -59,10 +141,27 @@ function parseWeights(bytes) {
         const [rows, cols = 1] = slice(name).shape;
         return new Tensor(rows, cols, f32(name), null, null);
     };
-    return {
+    const base = {
         embed: expand("embed"), sem: expand("sem"), importance: expand("importance"),
         w1: expand("w1"), b1: expand("b1"), w2: expand("w2"), b2: expand("b2"),
     };
+    if (!header["ap_w"])
+        return { ...base, tr: null };
+    // transformer encoder + attention pool
+    const nLayers = Object.keys(header).filter((k) => /^L\d+_qkv_w$/.test(k)).length;
+    const layers = [];
+    for (let l = 0; l < nLayers; l++) {
+        layers.push({
+            qkv_w: expand(`L${l}_qkv_w`), qkv_b: expand(`L${l}_qkv_b`),
+            o_w: expand(`L${l}_o_w`), o_b: expand(`L${l}_o_b`),
+            f1_w: expand(`L${l}_f1_w`), f1_b: expand(`L${l}_f1_b`),
+            f2_w: expand(`L${l}_f2_w`), f2_b: expand(`L${l}_f2_b`),
+            n1_w: expand(`L${l}_n1_w`), n1_b: expand(`L${l}_n1_b`),
+            n2_w: expand(`L${l}_n2_w`), n2_b: expand(`L${l}_n2_b`),
+        });
+    }
+    const tr = { layers, ap_w: expand("ap_w"), ap_b: expand("ap_b"), ap_q: expand("ap_q"), pn_w: expand("pn_w"), pn_b: expand("pn_b") };
+    return { ...base, tr };
 }
 /** Loads the bundled model and runs emoji inference. Construct once and reuse. */
 export class EmoModel {
@@ -103,21 +202,68 @@ export class EmoModel {
         if (ids.length === 0)
             ids = [this.meta.sem_pad_index];
         const sv = new Float32Array(semDim);
-        for (const id of ids) {
-            if (id >= sem.rows)
-                continue;
-            const base = id * semDim;
+        const tr = this.w.tr;
+        if (tr) {
+            // semantic sequence -> transformer encoder -> attention pool -> LayerNorm
+            const T = ids.length;
+            let S = new Float32Array(T * semDim);
+            for (let t = 0; t < T; t++) {
+                const id = ids[t] < sem.rows ? ids[t] : this.meta.sem_pad_index;
+                const sb = id * semDim, ob = t * semDim;
+                for (let c = 0; c < semDim; c++)
+                    S[ob + c] = id < sem.rows ? sem.get(sb + c) : 0;
+            }
+            const H = this.meta.n_heads ?? 4;
+            for (const L of tr.layers)
+                S = encoderLayer(S, T, semDim, H, L);
+            // attention pool: score_t = (aproj(tanh(S_t)) . q) / sqrt(d), softmax over t, weighted sum
+            const scoreScale = 1 / Math.sqrt(semDim);
+            const tanhS = new Float32Array(T * semDim);
+            for (let i = 0; i < T * semDim; i++)
+                tanhS[i] = Math.tanh(S[i]);
+            const proj = linear(tanhS, T, semDim, tr.ap_w, semDim, tr.ap_b);
+            const sc = new Float32Array(T);
+            let mx = -Infinity;
+            for (let t = 0; t < T; t++) {
+                let s = 0;
+                const pb = t * semDim;
+                for (let c = 0; c < semDim; c++)
+                    s += proj[pb + c] * tr.ap_q.get(c);
+                s *= scoreScale;
+                sc[t] = s;
+                if (s > mx)
+                    mx = s;
+            }
+            let sum = 0;
+            for (let t = 0; t < T; t++) {
+                const e = Math.exp(sc[t] - mx);
+                sc[t] = e;
+                sum += e;
+            }
+            for (let t = 0; t < T; t++) {
+                const a = sc[t] / sum, sb = t * semDim;
+                for (let c = 0; c < semDim; c++)
+                    sv[c] += a * S[sb + c];
+            }
+            layerNorm(sv, 1, semDim, tr.pn_w, tr.pn_b);
+        }
+        else {
+            for (const id of ids) {
+                if (id >= sem.rows)
+                    continue;
+                const base = id * semDim;
+                for (let c = 0; c < semDim; c++)
+                    sv[c] += sem.get(base + c);
+            }
+            for (let c = 0; c < semDim; c++)
+                sv[c] /= ids.length;
+            let norm = 0;
+            for (let c = 0; c < semDim; c++)
+                norm += sv[c] * sv[c];
+            norm = Math.sqrt(norm) + 1e-9;
             for (let c = 0; c < semDim; c++)
-                sv[c] += sem.get(base + c);
+                sv[c] /= norm;
         }
-        for (let c = 0; c < semDim; c++)
-            sv[c] /= ids.length;
-        let norm = 0;
-        for (let c = 0; c < semDim; c++)
-            norm += sv[c] * sv[c];
-        norm = Math.sqrt(norm) + 1e-9;
-        for (let c = 0; c < semDim; c++)
-            sv[c] /= norm;
         const inDim = ngDim + semDim;
         const x = new Float32Array(inDim);
         x.set(ng, 0);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@desert-ant-labs/emo",
-  "version": "0.5.0",
+  "version": "0.6.0",
   "description": "On-device emoji suggestions from text.",
   "type": "module",
   "main": "./dist/index.node.js",