diff --git a/DESIGN.md b/DESIGN.md index ce2f1a1..14ee30c 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -74,7 +74,7 @@ $$f_{\text{shell}}(n, \varepsilon) = 1 - (1-\varepsilon)^n$$ For any fixed $\varepsilon > 0$, $f_{\text{shell}} \to 1$ as $n \to \infty$. The shell thickness required to capture fraction $f$ is: -$$\varepsilon^{*}(f, n) = 1 - (1-f)^{1/n} \approx \frac{-\ln(1-f)}{n}$$ +$$\varepsilon^{\ast}(f, n) = 1 - (1-f)^{1/n} \approx \frac{-\ln(1-f)}{n}$$ | Fraction captured | Shell thickness | |:-----------------:|:---------------:| diff --git a/bun.lock b/bun.lock index e23cbb4..4ba4e33 100644 --- a/bun.lock +++ b/bun.lock @@ -5,12 +5,13 @@ "": { "name": "q2", "dependencies": { - "@huggingface/transformers": "4.0.0-next.8", + "@huggingface/transformers": "^4.0.0-next.8", }, "devDependencies": { "@eslint/eslintrc": "^3.3.5", "@html-eslint/eslint-plugin": "latest", "@html-eslint/parser": "latest", + "@playwright/test": "^1.58.2", "@types/node": "latest", "@typescript-eslint/eslint-plugin": "latest", "@typescript-eslint/parser": "latest", @@ -195,8 +196,6 @@ "@img/sharp-win32-x64": ["@img/sharp-win32-x64@0.34.5", "", { "os": "win32", "cpu": "x64" }, "sha512-+29YMsqY2/9eFEiW93eqWnuLcWcufowXewwSNIT6UwZdUUCrM3oFjMWH/Z6/TMmb4hlFenmfAVbpWeup2jryCw=="], - "@isaacs/fs-minipass": ["@isaacs/fs-minipass@4.0.1", "", { "dependencies": { "minipass": "^7.0.4" } }, "sha512-wgm9Ehl2jpeqP3zw/7mo3kRHFp5MEDhqAdwy1fTGkHAwnkGOVsgpvQhL8B5n1qlb01jV3n/bI0ZfZp5lWA1k4w=="], - "@istanbuljs/schema": ["@istanbuljs/schema@0.1.3", "", {}, "sha512-ZXRY4jNvVgSVQ8DL3LTcakaAtXwTVUxE81hslsyD2AtoXW/wVob10HkOJ1X/pAlcI7D+2YoZKg5do8G/w6RYgA=="], "@jridgewell/gen-mapping": ["@jridgewell/gen-mapping@0.3.13", "", { "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.0", "@jridgewell/trace-mapping": "^0.3.24" } }, "sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA=="], @@ -227,6 +226,8 @@ "@oxc-project/types": ["@oxc-project/types@0.115.0", "", {}, "sha512-4n91DKnebUS4yjUHl2g3/b2T+IUdCfmoZGhmwsovZCDaJSs+QkVAM+0AqqTxHSsHfeiMuueT75cZaZcT/m0pSw=="], + "@playwright/test": ["@playwright/test@1.58.2", "", { "dependencies": { "playwright": "1.58.2" }, "bin": { "playwright": "cli.js" } }, "sha512-akea+6bHYBBfA9uQqSYmlJXn61cTa+jbO87xVLCWbTqbWadRVmhxlXATaOjOgcBaWU4ePo0wB41KMFv3o35IXA=="], + "@polka/url": ["@polka/url@1.0.0-next.29", "", {}, "sha512-wwQAWhWSuHaag8c4q/KN/vCoeOJYshAIvMQwD4GpSb3OiZklFfvAgmj0VCBBImRpuF/aFgIRzllXlVX93Jevww=="], "@protobufjs/aspromise": ["@protobufjs/aspromise@1.1.2", "", {}, "sha512-j+gKExEuLmKwvz3OgROXtrJ2UG2x8Ch2YZUxahh+s1F2HZ+wAceUNLkvy6zKCPVRkU++ZWQrdxsUeQXmcg4uoQ=="], @@ -387,8 +388,6 @@ "chai": ["chai@6.2.2", "", {}, "sha512-NUPRluOfOiTKBKvWPtSD4PhFvWCqOi0BGStNWs57X9js7XGTprSmFoz5F0tWhR4WPjNeR9jXqdC7/UpSJTnlRg=="], - "chownr": ["chownr@3.0.0", "", {}, "sha512-+IxzY9BZOQd/XuYPRmrvEVjF/nqj5kgT4kEq7VofrDoM1MxoRjEWkrCC3EtLi59TVawxTAn+orJwFQcrqEN1+g=="], - "cli-cursor": ["cli-cursor@5.0.0", "", { "dependencies": { "restore-cursor": "^5.0.0" } }, "sha512-aCj4O5wKyszjMmDT4tZj93kxyydN/K5zPWSCe6/0AV/AA1pqe5ZBIw0a2ZfPQV7lL5/yb5HsUreJ6UFAF1tEQw=="], "cli-truncate": ["cli-truncate@5.2.0", "", { "dependencies": { "slice-ansi": "^8.0.0", "string-width": "^8.2.0" } }, "sha512-xRwvIOMGrfOAnM1JYtqQImuaNtDEv9v6oIYAs4LIHwTiKee8uwvIi363igssOC0O5U04i4AlENs79LQLu9tEMw=="], @@ -689,10 +688,6 @@ "minimatch": ["minimatch@3.1.5", "", { "dependencies": { "brace-expansion": "^1.1.7" } }, "sha512-VgjWUsnnT6n+NUk6eZq77zeFdpW2LWDzP6zFGrCbHXiYNul5Dzqk2HHQ5uFH2DNW5Xbp8+jVzaeNt94ssEEl4w=="], - "minipass": ["minipass@7.1.3", "", {}, "sha512-tEBHqDnIoM/1rXME1zgka9g6Q2lcoCkxHLuc7ODJ5BxbP5d4c2Z5cGgtXAku59200Cx7diuHTOYfSBD8n6mm8A=="], - - "minizlib": ["minizlib@3.1.0", "", { "dependencies": { "minipass": "^7.1.2" } }, "sha512-KZxYo1BUkWD2TVFLr0MQoM8vUUigWD3LlD83a/75BqC+4qE0Hb1Vo5v1FgcfaNXvfXzr+5EhQ6ing/CaBijTlw=="], - "mrmime": ["mrmime@2.0.1", "", {}, "sha512-Y3wQdFg2Va6etvQ5I82yUhGdsKrcYox6p7FfL1LbK2J4V01F9TGlepTIhnK24t7koZibmg82KGglhA1XK5IsLQ=="], "ms": ["ms@2.1.3", "", {}, "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA=="], @@ -839,8 +834,6 @@ "table": ["table@6.9.0", "", { "dependencies": { "ajv": "^8.0.1", "lodash.truncate": "^4.4.2", "slice-ansi": "^4.0.0", "string-width": "^4.2.3", "strip-ansi": "^6.0.1" } }, "sha512-9kY+CygyYM6j02t5YFHbNz2FN5QmYGv9zAjVp4lCDjlCw7amdckXlEt/bjMhUIfj4ThGRE4gCUH5+yGnNuPo5A=="], - "tar": ["tar@7.5.11", "", { "dependencies": { "@isaacs/fs-minipass": "^4.0.0", "chownr": "^3.0.0", "minipass": "^7.1.2", "minizlib": "^3.1.0", "yallist": "^5.0.0" } }, "sha512-ChjMH33/KetonMTAtpYdgUFr0tbz69Fp2v7zWxQfYZX4g5ZN2nOBXm1R2xyA+lMIKrLKIoKAwFj93jE/avX9cQ=="], - "tinybench": ["tinybench@2.9.0", "", {}, "sha512-0+DUvqWMValLmha6lr4kD8iAMK1HzV0/aKnCtWb9v9641TnP/MFb7Pc2bxoxQjTXAErryXVgUOfv2YqNllqGeg=="], "tinyexec": ["tinyexec@1.0.4", "", {}, "sha512-u9r3uZC0bdpGOXtlxUIdwf9pkmvhqJdrVCH9fapQtgy/OeTTMZ1nqH7agtvEfmGui6e1XxjcdrlxvxJvc3sMqw=="], @@ -919,7 +912,7 @@ "xmlchars": ["xmlchars@2.2.0", "", {}, "sha512-JZnDKK8B0RCDw84FNdDAIpZK+JuJw+s7Lz8nksI7SIuU3UXJJslUthsi+uWBUYOwPFwW7W7PRLRfUKpxjtjFCw=="], - "yallist": ["yallist@5.0.0", "", {}, "sha512-YgvUTfwqyc7UXVMrB+SImsVYSmTS8X/tSrtdNZMImM+n7+QTriRXyXim0mBrTXNeqzVF0KWGgHPeiyViFFrNDw=="], + "yallist": ["yallist@3.1.1", "", {}, "sha512-a4UGQaWPH59mOXUYnAG2ewncQS4i4F43Tv3JoAM+s2VDAmS9NsK8GpDMLrCHPksFT7h3K6TOoUNn2pb7RoXx4g=="], "yaml": ["yaml@2.8.2", "", { "bin": { "yaml": "bin.mjs" } }, "sha512-mplynKqc1C2hTVYxd0PU2xQAc22TI1vShAYGksCCfxbn/dFwnHTNi1bvYsBTkhdUNtGIf5xNOg938rrSSYvS9A=="], @@ -1005,8 +998,6 @@ "wrap-ansi/string-width": ["string-width@7.2.0", "", { "dependencies": { "emoji-regex": "^10.3.0", "get-east-asian-width": "^1.0.0", "strip-ansi": "^7.1.0" } }, "sha512-tsaTIkKW9b4N+AEj+SVA+WhJzV7/zMhcSu78mLKWSk7cXMOSHsBKFWUs0fWwq8QyK3MgJBQRX6Gbi4kYbdvGkQ=="], - "@babel/helper-compilation-targets/lru-cache/yallist": ["yallist@3.1.1", "", {}, "sha512-a4UGQaWPH59mOXUYnAG2ewncQS4i4F43Tv3JoAM+s2VDAmS9NsK8GpDMLrCHPksFT7h3K6TOoUNn2pb7RoXx4g=="], - "@eslint/config-array/minimatch/brace-expansion": ["brace-expansion@5.0.4", "", { "dependencies": { "balanced-match": "^4.0.2" } }, "sha512-h+DEnpVvxmfVefa4jFbCf5HdH5YMDXRsmKflpf1pILZWRFlTbJpxeU55nJl4Smt5HQaGzg1o6RHFPJaOqnmBDg=="], "@typescript-eslint/typescript-estree/minimatch/brace-expansion": ["brace-expansion@5.0.4", "", { "dependencies": { "balanced-match": "^4.0.2" } }, "sha512-h+DEnpVvxmfVefa4jFbCf5HdH5YMDXRsmKflpf1pILZWRFlTbJpxeU55nJl4Smt5HQaGzg1o6RHFPJaOqnmBDg=="], diff --git a/src/app.ts b/src/app.ts index 8ed5b44..19814eb 100644 --- a/src/app.ts +++ b/src/app.ts @@ -25,16 +25,13 @@ import type { ChatMessage, GenerationConfig, EmbeddingMsg, + ModelOutputsMsg, + Q2Msg, } from './types.js'; import { - getKernel, l2Normalise, q2EncodeDirect, q2KeyDirect, - DTYPE_TO_Q2, - Q2_DTYPE_FP32, - Q2_INPUT_OFFSET, - Q2_OUTPUT_OFFSET, } from './q2.js'; import { deleteStoredFile, @@ -681,6 +678,12 @@ export function handleWorkerMessage(msg: WorkerOutMsg): void { case 'embedding': onEmbedding(msg); break; + case 'model-outputs': + onModelOutputs(msg); + break; + case 'q2': + onQ2(msg); + break; case 'done': onDone(); break; @@ -828,55 +831,48 @@ export function onEmbedding(msg: EmbeddingMsg): void { `Shape: [${seqLen} × ${hiddenDim}] dtype=${dtype} stats=unavailable`; } - // ── Q² kernel ──────────────────────────────────────────────────────────── - // Run the quaternary quantisation in the background. The WASM kernel is - // preferred; if instantiation fails (e.g. in test environments that lack - // WebAssembly.instantiate) we fall back to the pure-TS implementation. - const n = hiddenDim; - const dtypeId = DTYPE_TO_Q2[dtype] ?? Q2_DTYPE_FP32; +} - if (seqLen < 1) { - appLog('warn', 'Q² embedding: seqLen < 1; skipping quantisation', { seqLen }); - return; +/** + * Handles the compact Q² quantisation result sent by the worker kernel. + * + * The worker runs the Q² WASM kernel before sending, so only packed bytes + * and the 64-bit key cross the thread boundary (see worker.ts quantiseAndSend). + */ +/** + * Shows the user which ONNX output nodes the loaded model exports and whether + * Q² fingerprinting was able to locate a hidden-state tensor among them. + * + * Called once per generation turn, immediately after the embedding forward + * pass in the worker. Surfaced in the embedding panel so the user knows + * exactly why Q² may be unavailable and what the model actually exports. + */ +export function onModelOutputs(msg: ModelOutputsMsg): void { + appLog('info', 'onModelOutputs received', msg); + embeddingPanel.classList.remove('hidden'); + + // Format each output as name[d0×d1×…] for compact display. + const outputList = Object.entries(msg.outputs) + .map(([name, dims]) => `${name}[${dims.join('×')}]`) + .join(' '); + + if (msg.hiddenStateKey !== null) { + embeddingStats.textContent = + `ONNX outputs: ${outputList}\n` + + `Q² using: ${msg.hiddenStateKey}[${(msg.outputs[msg.hiddenStateKey] ?? []).join('×')}]`; + } else { + embeddingStats.textContent = + `ONNX outputs: ${outputList}\n` + + `Q² unavailable — no 3-D hidden-state output found.\n` + + `To enable Q² fingerprinting, re-export the model with a last_hidden_state ` + + `(or equivalent) output node, or use a model that already exports one.`; } +} - appLog('debug', 'onEmbedding: starting Q² kernel', { hiddenDim: n, dtypeId, seqLen }); - void (async () => { - try { - const kernel = await getKernel(); - const mem = new Uint8Array(kernel.memory.buffer); - - // Copy the raw activation buffer into WASM memory at the input offset. - const inputBytes = new Uint8Array(msg.data); - mem.set(inputBytes, Q2_INPUT_OFFSET); - - // Run quantisation: L2-normalise last token, threshold, Gray-encode. - kernel.quantise(Q2_INPUT_OFFSET, seqLen, n, dtypeId, Q2_OUTPUT_OFFSET); - - // Derive the 64-bit transition key. - const rawKey = kernel.key(Q2_OUTPUT_OFFSET, n); - const key = BigInt.asUintN(64, rawKey); - - appLog('debug', 'Q² WASM kernel produced key', { key: `0x${key.toString(16).padStart(16, '0')}`, hiddenDim: n }); - // Read back packed bytes. - const packed = new Uint8Array(kernel.memory.buffer, Q2_OUTPUT_OFFSET, n >> 2); - renderQ2Result(packed, key, n, currentSettings.q2KeyDisplayMode); - } catch { - // WASM unavailable — use the pure-TypeScript fallback (fp32 only). - // This path is taken in test environments and SSR contexts. - // For sub-fp32 dtypes the WASM kernel is required; log a warning and skip. - if (dtype !== 'fp32') { - appLog('warn', 'Q² TS fallback: non-fp32 dtype requires WASM kernel; skipping', { dtype }); - return; - } - appLog('debug', 'Q² falling back to TS implementation', { seqLen, hiddenDim: n }); - const all = new Float32Array(msg.data); - const vec = l2Normalise(all.subarray((seqLen - 1) * n, seqLen * n), n); - const { packed, key } = q2EncodeDirect(vec, n); - appLog('debug', 'Q² TS fallback produced key', { key: `0x${BigInt.asUintN(64, key).toString(16).padStart(16, '0')}`, hiddenDim: n }); - renderQ2Result(packed, BigInt.asUintN(64, key), n, currentSettings.q2KeyDisplayMode); - } - })(); +export function onQ2(msg: Q2Msg): void { + const packed = new Uint8Array(msg.packed); + appLog('debug', 'onQ2 received', { n: msg.n, key: `0x${msg.key.toString(16).padStart(16, '0')}` }); + renderQ2Result(packed, msg.key, msg.n, currentSettings.q2KeyDisplayMode); } export function onDone(): void { diff --git a/src/types.ts b/src/types.ts index cea515f..58b6aec 100644 --- a/src/types.ts +++ b/src/types.ts @@ -101,6 +101,47 @@ export interface EmbeddingMsg { dtype: 'fp32' | 'fp16' | 'q8' | 'q4' | 'q2'; } +/** + * Sent once per generation turn immediately after the embedding forward pass, + * regardless of whether a usable hidden-state output was found. + * + * Lets the main thread show the user exactly which ONNX output nodes the + * loaded model exposes and explain why Q² fingerprinting may be unavailable. + */ +export interface ModelOutputsMsg { + type: 'model-outputs'; + /** + * Every output node the model's ONNX session exposes. + * Key: node name. Value: dimension array, e.g. [1, 42, 4096]. + */ + outputs: Record; + /** + * The output node name that was selected for Q² quantisation, + * or null when no suitable hidden-state tensor was found. + */ + hiddenStateKey: string | null; +} + +/** + * Q² quantisation result produced by the worker kernel. + * + * The worker runs the Q² WASM kernel immediately after extracting an embedding, + * so only the compact quantised representation crosses the thread boundary + * instead of the raw activation buffer (~64× smaller for fp32 n=4096). + */ +export interface Q2Msg { + type: 'q2'; + /** + * n/4 packed Gray-encoded bytes (transferable ArrayBuffer). + * Transfer via postMessage(msg, [packed]) to avoid structured-clone copy. + */ + packed: ArrayBuffer; + /** 64-bit MSB-aligned transition key (DESIGN.md §2.2). */ + key: bigint; + /** Original embedding dimension (n). */ + n: number; +} + export interface DoneMsg { type: 'done'; } @@ -115,6 +156,8 @@ export type WorkerOutMsg = | ProgressMsg | TokenMsg | EmbeddingMsg + | ModelOutputsMsg + | Q2Msg | DoneMsg | ErrorMsg; diff --git a/src/worker.ts b/src/worker.ts index fcfe999..1107337 100644 --- a/src/worker.ts +++ b/src/worker.ts @@ -38,7 +38,15 @@ import type { GenerationConfig, EmbeddingMsg, Dtype, + ModelOutputsMsg, } from './types.js'; +import { + getKernel, + DTYPE_TO_Q2, + Q2_DTYPE_FP32, + Q2_INPUT_OFFSET, + Q2_OUTPUT_OFFSET, +} from './q2.js'; /** Subset of the ProgressInfo union we care about for download tracking. */ interface DownloadProgress { @@ -253,6 +261,60 @@ function send(msg: WorkerOutMsg, transfer: Transferable[] = []): void { workerScope.postMessage(msg, transfer); } +/** + * Run the Q² WASM kernel on a raw embedding buffer and send the compact result + * to the main thread. + * + * By quantising in the worker we avoid transferring the full activation buffer + * across the thread boundary. Only the packed output (n/4 bytes) and the + * 64-bit key are sent — roughly 64× less data than the raw fp32 input for a + * typical hidden dimension of 4096. + * + * @param embeddingBuffer - raw activation bytes (owned; will NOT be transferred) + * @param seqLen - number of token positions in the buffer + * @param hiddenDim - embedding dimension n + * @param dtype - element dtype of the activation buffer + */ +async function quantiseAndSend( + embeddingBuffer: ArrayBuffer, + seqLen: number, + hiddenDim: number, + dtype: EmbeddingMsg['dtype'], +): Promise { + if (seqLen < 1) { + workerLog('warn', 'quantiseAndSend: seqLen < 1; skipping Q² quantisation', { seqLen }); + return; + } + if (hiddenDim < 1) { + workerLog('warn', 'quantiseAndSend: hiddenDim < 1; skipping Q² quantisation', { hiddenDim }); + return; + } + const n = hiddenDim; + const dtypeId = DTYPE_TO_Q2[dtype] ?? Q2_DTYPE_FP32; + try { + const kernel = await getKernel(); + const mem = new Uint8Array(kernel.memory.buffer); + + // Copy activation bytes into WASM linear memory at the fixed input offset. + mem.set(new Uint8Array(embeddingBuffer), Q2_INPUT_OFFSET); + + // L2-normalise last token, threshold, Gray-encode → packed output at Q2_OUTPUT_OFFSET. + kernel.quantise(Q2_INPUT_OFFSET, seqLen, n, dtypeId, Q2_OUTPUT_OFFSET); + + // Derive the 64-bit transition key. + const key = BigInt.asUintN(64, kernel.key(Q2_OUTPUT_OFFSET, n)); + + // Slice to an independent buffer so we can transfer ownership without + // detaching the WASM module's shared memory view. + const packed = new Uint8Array(kernel.memory.buffer, Q2_OUTPUT_OFFSET, n >> 2).slice(); + + workerLog('debug', 'Q² kernel produced key', { key: `0x${key.toString(16).padStart(16, '0')}`, n }); + send({ type: 'q2', packed: packed.buffer, key, n }, [packed.buffer]); + } catch (err) { + workerLog('warn', 'Q² kernel failed; skipping quantisation result', { error: err }); + } +} + // ─── Model loading ──────────────────────────────────────────────────────────── /** @@ -450,27 +512,112 @@ async function generateResponse( outputLength: output.length, }); - // ── Embedding extraction ───────────────────────────────────────────────── - // NOTE: Accessing per-layer hidden states during generation is NOT - // supported by the transformers.js text-generation pipeline. The - // model.generate() loop in transformers.js v3 does not collect hidden - // states — output_hidden_states in GenerationConfig has no effect. + // ── Embedding extraction via direct model forward pass ─────────────────── + // The text-generation pipeline returns decoded text, not raw tensors. + // To get the last-token hidden state we tokenize the full generated + // conversation and call pipe.model() directly (one additional forward pass, + // no KV cache, O(seqLen) attention). // - // The correct approach to obtain hidden states is: - // 1. Use a feature-extraction pipeline with a dedicated embedding model. - // 2. Or call pipe.model.forward() on the generated token sequence with - // a model that exports intermediate hidden states in its ONNX graph. + // Hidden-state detection strategy (in order): + // 1. Try each name in HIDDEN_STATE_CANDIDATES. + // 2. Fall back to any 3-D output (shape [batch, seq, hidden]) — this + // handles models like LFM2.5 that use non-standard output names. // - // Standard onnx-community text-generation models export only - // {logits, past_key_values}. To use Q² fingerprinting, configure a - // dedicated embedding model via the benchModelT3 setting. - const extConfig = config as GenerationConfig & { return_embeddings?: boolean }; - const wantEmbeddings = extConfig.return_embeddings === true; - if (wantEmbeddings) { - workerLog('warn', - 'Embedding extraction via text-generation pipeline is not supported. ' + - 'Use a feature-extraction pipeline with a dedicated embedding model instead.'); - } + // Regardless of outcome we send a ModelOutputsMsg so the main thread + // can show the user exactly what the model exports and why Q² may be + // unavailable. + void (async () => { + try { + // Reconstruct the full conversation text from the pipeline output so we + // can re-tokenize it for the embedding forward pass. + const fullConv = output[0]?.generated_text; + const convText = + typeof fullConv === 'string' + ? fullConv + : (fullConv as ChatMessage[] ?? []).map((m: ChatMessage) => m.content).join('\n'); + + // Tokenize without padding — we want the exact sequence length so the + // last token position maps cleanly to the final hidden-state row. + const tokenized = (pipe!.tokenizer as unknown as ( + text: string, + opts: Record, + ) => Record)(convText, { + return_tensors: 'pt', + truncation: true, + }); + + // Direct model call. In transformers.js the PreTrainedModel is callable; + // it runs the underlying ONNX session and returns an object whose keys + // are the ONNX graph output names. + type OnnxTensor = { dims: number[]; data: Float32Array }; + const modelCallable = pipe!.model as unknown as ( + inputs: Record, + ) => Promise>; + const modelOutput = await modelCallable(tokenized); + + // Build the shape map for ModelOutputsMsg (sent regardless of outcome). + const outputShapes: Record = {}; + for (const [k, v] of Object.entries(modelOutput)) { + if (v?.dims) outputShapes[k] = v.dims; + } + workerLog('info', 'Model ONNX outputs', outputShapes); + + // ── Locate the hidden-state tensor ────────────────────────────────── + // Well-known names tried first, then any 3-D output as a catch-all so + // non-standard architectures (e.g. LFM2.5) are detected automatically. + const HIDDEN_STATE_CANDIDATES = [ + 'last_hidden_state', // standard HF / transformers.js + 'hidden_states', + 'last_conv_hidden_states', // possible LFM2.5 naming + 'encoder_last_hidden_state', + ]; + + let hiddenStateKey: string | null = null; + let hiddenState: OnnxTensor | null = null; + + for (const key of HIDDEN_STATE_CANDIDATES) { + const t = modelOutput[key]; + if (t?.dims.length === 3) { hiddenStateKey = key; hiddenState = t; break; } + } + if (!hiddenState) { + // Shape-based fallback: any 3-D output is likely a hidden state. + for (const [key, t] of Object.entries(modelOutput)) { + if (t?.dims.length === 3) { + hiddenStateKey = key; + hiddenState = t; + workerLog('info', + `Hidden state auto-detected via shape: "${key}" dims=${JSON.stringify(t.dims)}`); + break; + } + } + } + + // Always inform the main thread of what the model exported. + const modelOutputsMsg: ModelOutputsMsg = { + type: 'model-outputs', + outputs: outputShapes, + hiddenStateKey, + }; + send(modelOutputsMsg); + + if (!hiddenState || hiddenStateKey === null) { + workerLog('warn', + 'No 3-D hidden-state output found; Q² fingerprinting unavailable for this model.', + { availableOutputs: Object.keys(outputShapes) }); + return; + } + + const [, seqLen, hiddenDim] = hiddenState.dims; + // hiddenState.data is a shared view backed by the ONNX runtime buffer; + // slice() makes an independent copy safe to hand off to the Q² kernel. + const data = hiddenState.data.slice().buffer; + workerLog('info', 'Embedding forward pass complete', + { key: hiddenStateKey, seqLen, hiddenDim }); + await quantiseAndSend(data, seqLen!, hiddenDim!, 'fp32'); + } catch (err) { + workerLog('warn', 'Embedding extraction failed', { error: err }); + } + })(); send({ type: 'done' }); } catch (err) { diff --git a/test/app.test.ts b/test/app.test.ts index b62695a..7ab188a 100644 --- a/test/app.test.ts +++ b/test/app.test.ts @@ -249,6 +249,18 @@ describe('app.ts helpers and DOM integration', () => { expect(stats.textContent).toContain('dtype=fp32'); }); + it('onQ2 renders the Q² result without needing the raw embedding buffer', () => { + const stats = document.querySelector('#embedding-stats') as HTMLElement; + stats.textContent = ''; + + // n=8 → 2 packed bytes; key produced by the worker kernel + const packed = new Uint8Array([0xAA, 0xAA]); // D D D D D D D D + app.onQ2({ type: 'q2', packed: packed.buffer, key: 0xdd8c000000000000n, n: 8 }); + + expect(stats.textContent).toContain('Q²:'); + expect(stats.textContent).toContain('2 bytes'); + }); + it('sendMessage posts a generate message and updates the UI', async () => { const input = document.querySelector('#user-input') as HTMLTextAreaElement; const sendBtn = document.querySelector('#send-btn') as HTMLButtonElement;