@atlaskit/editor-plugin-autocomplete 3.0.0 → 3.2.0

package/src/pm-plugins/local-slow-lane-client.ts

@@ -1,18 +1,27 @@
 /**
  * Local Slow Lane Client: On-device inference via @mlc-ai/web-llm.
  *
- * Drop-in replacement for the network-based slow-lane-client.  Instead of
- * calling a backend API, this client uses MLC WebLLM to run a small language
- * model (SmolLM 135M) directly in the browser via WebGPU.
+ * Drop-in replacement for the network-based slow-lane-client. Instead of calling
+ * a backend API, this client runs two models in the browser via WebGPU, in a
+ * single MLCEngine, to reproduce the BE encoder's outputs on-device:
+ *
+ *   - Causal LM (SmolLM2-135M-Instruct): one decode step per word boundary. A
+ *     registered LogitProcessor captures the raw next-token logits, which
+ *     `computeBePayload` turns into a whole-word `lm_logits` payload — a faithful
+ *     port of the BE `CausalLMEncoder._get_top_k_probs` (masked softmax over the
+ *     vocab's first-tokens, prefix expansion, L2 reservation, log-space pooling).
+ *   - Semantic embedder (Snowflake Arctic Embed S): produces the real 384-d
+ *     `semantic_vector`. Inputs are wrapped as passages (see `wrapForArctic`) so
+ *     the runtime vector lands in the same space as the precomputed word bin.
  *
  * ── Why main thread (no Web Worker)? ─────────────────────────────────────
- * SmolLM 135M is small enough (~270 MB weights, 350-400 MB VRAM) that
- * WebGPU inference on the main thread is production-viable:
+ * The models are small enough (~640 MB combined VRAM) that WebGPU inference on
+ * the main thread is viable:
  *
  *   - WebGPU GPU compute is inherently async (doesn't block the main thread)
- *   - CPU overhead (tokenization + post-processing) is only 5-10 ms
- *   - Single forward pass latency is 50-150 ms — well within autocomplete
- *     expectations (~250 ms between word boundaries)
+ *   - CPU overhead (BE-parity post-processing) is a few ms
+ *   - Per-inference latency is well within autocomplete expectations
+ *     (~250 ms between word boundaries)
  *
  * This avoids all the complexity of Web Workers:
  *   - No CSP workarounds (blob URLs, inline scripts)
@@ -26,7 +35,7 @@
  * asynchronously after each updateContext() call.
  */
 
-import type { MLCEngine, InitProgressReport, AppConfig } from '@mlc-ai/web-llm';
+import type { MLCEngine, InitProgressReport, AppConfig, LogitProcessor } from '@mlc-ai/web-llm';
 
 import { isAutocompleteDebugEnabled } from './debug-mode';
 import { isWordBoundary } from './slow-lane-client';
@@ -85,20 +94,437 @@ export interface LocalSlowLaneClient {
 
 const DEFAULT_DEBOUNCE_MS = 300;
 
-export const LOCAL_MLC_MODEL_ID = 'SmolLM2-135M-Instruct-q0f16-MLC';
+export const LOCAL_MLC_CAUSAL_MODEL_ID = 'SmolLM2-135M-Instruct-q0f16-MLC';
+
+/**
+ * MLC ID for the semantic embedder (Snowflake Arctic Embed S, batch=4 variant).
+ *
+ * The `-b4` suffix selects the prebuilt variant compiled for a max batch size of
+ * 4 (≈239 MB VRAM) rather than `-b32` (≈1023 MB VRAM). Autocomplete embeds one
+ * context at a time, so `-b4` is the right fit. This model IS in
+ * `prebuiltAppConfig.model_list` of web-llm 0.2.82 — no `customModelConfig` needed.
+ */
+export const LOCAL_MLC_EMBEDDING_MODEL_ID = 'snowflake-arctic-embed-s-q0f32-MLC-b4';
+
+/**
+ * Wrap raw context text with BERT special tokens before embedding.
+ *
+ * web-llm's `EmbeddingPipeline` does NOT auto-prepend `[CLS]` / append `[SEP]`
+ * (the official MLC embeddings example wraps manually). The Python
+ * `sentence_transformers` side that generated the word-vector bin adds these
+ * inside `model.encode()`, so we must mirror it here for the runtime context
+ * vector to land in the same region of Arctic's embedding space as the bin.
+ *
+ * No query prefix is applied: the semantic step is sentence-to-sentence (`s2s`)
+ * similarity ("which words are conceptually similar to this context?"), not
+ * sentence-to-passage (`s2p`) retrieval. Arctic's query prefix would misframe
+ * the relationship. Encode both sides as passages. See implementation.md §4.3.
+ */
+export const wrapForArctic = (text: string): string => `[CLS] ${text} [SEP]`;
+
+/**
+ * BE-parity constants — must match `CausalLMEncoder` defaults in the Python
+ * sidecar (`cc-smarts/python-sidecar/src/causal_lm_encoder.py`) and
+ * `SlowLaneEngine` (`typeahead_context_encoding.py`) so local payloads behave
+ * identically to the server-client setup.
+ */
+export const BE_PARITY = {
+	/** Final payload size cap (BE: `top_k_words`). */
+	TOP_K_WORDS: 2000,
+	/** L2 (domain) words admitted unconditionally before pooling (BE: `reserved_l2_slots`). */
+	RESERVED_L2_SLOTS: 500,
+	/** Log-space additive bias favouring L2 over L3 in the pool (BE: `l2_bias`). */
+	L2_BIAS: 1.0,
+	/** Drop words below this probability from the final payload (BE: `> 0.00001`). */
+	MIN_PROB: 0.00001,
+	/**
+	 * Word-level approximation of the BE causal LM token limit.
+	 *
+	 * BE: `CausalLMEncoder.max_context_tokens = 100` (BPE tokens, left-truncated).
+	 * FE: no tokenizer available, so we approximate with word count. English text
+	 * averages ~1.3–1.5 BPE tokens/word, meaning 100 words ≈ 130–150 tokens.
+	 * Using 100 words keeps the approximation simple and errs on the side of
+	 * sending slightly more context than the BE sees — acceptable for a PoC.
+	 */
+	MAX_CONTEXT_TOKENS: 100,
+	/**
+	 * Word-level rolling window for the semantic embedder.
+	 *
+	 * BE: `SlowLaneEngine.max_context_words = 100` (applied in
+	 * `typeahead_context_encoding.py` before calling `SemanticEncoder.encode`).
+	 * Truncated identically here so the runtime Arctic vector lands in the same
+	 * region of the embedding space as the precomputed word-vector bin.
+	 */
+	MAX_CONTEXT_WORDS: 100,
+} as const;
+
+/**
+ * Return the last `n` whitespace-separated words of `text`, joined by spaces.
+ * Mirrors the BE rolling-window truncation applied before both encoders.
+ */
+const truncateToLastNWords = (text: string, n: number): string => {
+	const words = text.trim().split(/\s+/u);
+	return words.length <= n ? text : words.slice(-n).join(' ');
+};
+
+// ─── Logit capture ─────────────────────────────────────────────────────────
 
-/** HF root for the default weights (includes `tensor-cache.json` for WebLLM 0.2+). */
-export const LOCAL_MLC_HF_MODEL_REPO =
-	'https://huggingface.co/mlc-ai/SmolLM2-135M-Instruct-q0f16-MLC';
+/**
+ * A LogitProcessor that captures the raw next-token logits and passes them
+ * through unmodified.
+ *
+ * web-llm invokes `processLogits` on the CPU after the model's forward pass and
+ * before sampling, handing us the full `Float32Array(vocab_size)` at the current
+ * decode position. We copy it off web-llm's shared buffer (which it may reuse
+ * across calls) and return the original untouched so sampling is unaffected.
+ *
+ * This is the raw-logit access the BE-parity algorithm needs (masked softmax +
+ * prefix expansion, consumed in a later step). Registered for the causal LM
+ * only — the embedder never decodes tokens, so it produces no logits.
+ */
+class CapturingLogitProcessor implements LogitProcessor {
+	captured: Float32Array | null = null;
 
-export const LOCAL_MLC_MODEL_LIB_WASM_NAME = 'SmolLM2-135M-Instruct-q0f16-ctx4k_cs1k-webgpu.wasm';
+	processLogits = (logits: Float32Array): Float32Array => {
+		// Copy off web-llm's shared buffer — it may reuse `logits` across calls.
+		this.captured = new Float32Array(logits);
+		return logits;
+	};
+
+	processSampledToken = (): void => {
+		// No-op — we don't track sampled tokens.
+	};
+
+	resetState = (): void => {
+		this.captured = null;
+	};
+}
+
+// ─── BE-parity data + algorithm ──────────────────────────────────────────────
 
 /**
- * Original target repo (add-basics fine-tune). **Not compatible with WebLLM 0.2.x** (no `tensor-cache.json`).
- * @see module doc above
+ * Prefix-expansion map: first-token id → words whose space-prefixed SmolLM2
+ * encoding starts with that token. Generated offline by
+ * `scripts/gen_first_token_to_words.py`, which mirrors the BE's in-memory map
+ * (`CausalLMEncoder._ensure_loaded`).
+ *
+ * Populated lazily by `loadBePayloadData()` from a dynamically-imported JSON so
+ * the (large) payload is only fetched when the local client is actually
+ * initialised — keeping it out of the editor's main chunk for the vast majority
+ * of users (who run with `useLocalModel` off).
  */
-export const HUGGINGFACE_TB_SMOLLM_ADD_BASICS_REPO =
-	'https://huggingface.co/HuggingFaceTB/smollm-135M-instruct-add-basics-q0f16-MLC';
+let firstTokenToWords: Map<number, string[]> = new Map();
+
+/**
+ * L2 (Atlassian-domain) word set, derived from the keys of `vocabulary_10k.json`.
+ * Used by `computeBePayload` for tier-aware ranking: any word in the prefix map
+ * that is not in this set is treated as L3 (general English), matching the BE.
+ * Populated lazily alongside `firstTokenToWords` — see `loadBePayloadData()`.
+ */
+let l2Words: Set<string> = new Set();
+
+/**
+ * Array of token IDs that appear as a first token for at least one vocabulary
+ * word. Derived from `firstTokenToWords` when the data loads so `computeBePayload`
+ * does not re-allocate this array on every word-boundary call.
+ */
+let prefixMapTokenIds: number[] = [];
+
+/** De-dupes concurrent loads and lets repeated calls await the same payload. */
+let bePayloadDataPromise: Promise<void> | undefined;
+
+/**
+ * Unwrap a dynamically imported JSON module to the parsed JSON value, working
+ * across the two interop modes AFM's bundler chain emits:
+ *
+ *   1. **`.default`-wrapped namespace** — classic webpack (and Jest) hang the
+ *      JSON value under the `default` export.
+ *   2. **Named-exports namespace** — webpack 5 / atlaspack with JSON
+ *      named-exports (or native ESM JSON modules) expose each top-level key as
+ *      a named export and shadow `default`, so `mod.default` can be `undefined`
+ *      (or some unrelated value) even though `mod` itself holds the data.
+ *
+ * The caller MUST declare the underlying JSON shape via `shape` because, in
+ * named-exports mode, a dense array `["a","b"]` and a sparse numeric-keyed
+ * object `{"5":"a","12":"b"}` are emitted identically (`{"0":..}` / `{"5":..}`);
+ * no runtime heuristic can tell them apart, so only the caller knows which:
+ *
+ *   - `'object'` — the JSON is a `{...}` (including sparse maps keyed by integer
+ *     IDs). The named exports are rebuilt into a plain object so `Object.entries`
+ *     yields the real keys, not synthetic array indices.
+ *   - `'array'` — the JSON is a `[...]`, reconstructed from the `0..n-1` indices.
+ *
+ * :param mod: The raw module object returned by `await import('./*.json')`.
+ * :param shape: `'object'` if the source JSON is `{...}`, `'array'` if `[...]`.
+ * :returns: The parsed JSON value, or `null` if neither interop mode applies.
+ */
+const unwrapJsonModule = <T,>(mod: unknown, shape: 'object' | 'array'): T | null => {
+	if (mod == null || typeof mod !== 'object') {
+		return null;
+	}
+	const namespace = mod as Record<string, unknown> & { default?: unknown };
+
+	// Compute the named-export own-keys (strip synthetic markers).
+	const ownKeys = Object.keys(namespace).filter(
+		(k) => k !== 'default' && k !== '__esModule',
+	);
+
+	// PREFER named exports when present — they always reflect the JSON's real
+	// top-level keys / indices, regardless of what `default` happens to be.
+	// Under JSON named-exports mode `default` is not necessarily the parsed
+	// value (e.g. for `{"service": 0, ...}` it can be the number `0`, with the
+	// real data in the named exports), so taking `default` first would corrupt it.
+	if (ownKeys.length > 0) {
+		if (shape === 'array') {
+			// JSON arrays are dense; reconstruct from `0..length-1` indices.
+			const len = ownKeys.length;
+			const arr = new Array(len);
+			for (let i = 0; i < len; i++) {
+				arr[i] = namespace[String(i)];
+			}
+			return arr as T;
+		}
+		// shape === 'object'. Rebuild a plain object from the (stripped) own
+		// keys so callers can `Object.entries()` it without iterating over
+		// `default` / `__esModule`, and to detach from the module-namespace
+		// object (which is sealed/non-extensible on some bundler outputs).
+		const obj: Record<string, unknown> = {};
+		for (const k of ownKeys) {
+			obj[k] = namespace[k];
+		}
+		return obj as T;
+	}
+
+	// Fallback: no named exports — classic webpack JSON-module interop where
+	// the whole parsed JSON value is hung under `default`. Trust it.
+	if ('default' in namespace && namespace.default != null) {
+		return namespace.default as T;
+	}
+
+	return null;
+};
+
+/**
+ * Lazily load and build the BE-parity lookup tables from their JSON payloads.
+ * The dynamic imports are split into their own async chunks so neither file is
+ * bundled into the editor's main chunk unless local inference is initialised.
+ *
+ * :returns:
+ *   A promise that resolves once `firstTokenToWords`, `l2Words` and
+ *   `prefixMapTokenIds` are populated.
+ */
+const loadBePayloadData = (): Promise<void> => {
+	if (!bePayloadDataPromise) {
+		bePayloadDataPromise = (async () => {
+			const [firstTokenToWordsModule, vocabularyModule] = await Promise.all([
+				import(
+					/* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-first-token-to-words" */ './data/first_token_to_words.json'
+				),
+				import(
+					/* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-vocabulary-10k" */ './data/vocabulary_10k.json'
+				),
+			]);
+
+			const firstTokenToWordsData = unwrapJsonModule<Record<string, string[]>>(
+				firstTokenToWordsModule,
+				'object',
+			);
+			const vocabularyData = unwrapJsonModule<{ words: Record<string, unknown> }>(
+				vocabularyModule,
+				'object',
+			);
+
+			if (firstTokenToWordsData == null || vocabularyData?.words == null) {
+				// Hard-fail with a precise message so the catch() in initEngine logs
+				// exactly which import couldn't be unwrapped, rather than the generic
+				// V8 "Cannot convert undefined or null to object" we hit before the
+				// helper was added.
+				throw new Error(
+					`[LocalSlowLane] JSON module could not be unwrapped — ` +
+						`firstTokenToWordsData=${firstTokenToWordsData == null ? 'null/undefined' : 'defined'}, ` +
+						`vocabularyData=${vocabularyData == null ? 'null/undefined' : vocabularyData.words == null ? 'defined but missing .words' : 'defined'}`,
+				);
+			}
+
+			firstTokenToWords = new Map(
+				Object.entries(firstTokenToWordsData).map(([tokenId, words]) => [
+					Number(tokenId),
+					words,
+				]),
+			);
+			l2Words = new Set(Object.keys(vocabularyData.words));
+			prefixMapTokenIds = Array.from(firstTokenToWords.keys());
+
+			if (isAutocompleteDebugEnabled()) {
+				// eslint-disable-next-line no-console
+				console.log(
+					'%c[LocalSlowLane] %c✅ BE-parity payload data loaded:',
+					'color: #9c27b0; font-weight: bold;',
+					'color: #4caf50; font-weight: bold;',
+					{
+						firstTokenToWordsEntries: firstTokenToWords.size,
+						l2WordsCount: l2Words.size,
+						prefixMapTokenIdsLength: prefixMapTokenIds.length,
+					},
+				);
+			}
+		})().catch((e) => {
+			// Don't cache a rejected promise — a transient import failure would
+			// otherwise prevent the local model from ever initialising again this
+			// session. Reset so the next init attempt retries.
+			bePayloadDataPromise = undefined;
+			throw e;
+		});
+	}
+	return bePayloadDataPromise;
+};
+
+/**
+ * Convert a raw next-token logit vector into a whole-word probability payload,
+ * faithfully porting the BE `CausalLMEncoder._get_top_k_probs`
+ * (`cc-smarts/python-sidecar/src/causal_lm_encoder.py`).
+ *
+ * Steps: (1) numerically-stable masked softmax over only the token ids present
+ * in the prefix-expansion map; (2) spread each token's probability to every
+ * whole word sharing that first token, taking the max; (3) reserve the top L2
+ * words unconditionally; (4) rank the remainder in a log-space pool with an
+ * additive L2 bias; (5) emit raw probabilities for the survivors, lowercased
+ * and trimmed at `MIN_PROB`.
+ *
+ * :params:
+ *   rawLogits: Full-vocabulary logits from the LM's single decode step
+ *   prefixMap: Map of first-token id to the words starting with that token
+ *   domainWords: Set of L2 (domain) words, for tier-aware ranking
+ * :returns:
+ *   A record of lowercase word to probability — the BE `lm_logits` payload
+ */
+export const computeBePayload = (
+	rawLogits: Float32Array,
+	prefixMap: Map<number, string[]>,
+	domainWords: Set<string>,
+	/**
+	 * Pre-derived token-ID array for the softmax mask. Defaults to the
+	 * module-level `prefixMapTokenIds` (zero allocation in production). Pass
+	 * `Array.from(prefixMap.keys())` in tests that supply a custom prefixMap so
+	 * the softmax mask stays consistent with the iteration in Step 2.
+	 */
+	validTokenIds: number[] = prefixMapTokenIds,
+): Record<string, number> => {
+
+	// 1. Numerically-stable masked softmax over validTokenIds only.
+	let maxLogit = -Infinity;
+	for (const id of validTokenIds) {
+		const v = rawLogits[id];
+		if (v > maxLogit) {
+			maxLogit = v;
+		}
+	}
+	let sumExp = 0;
+	const expByToken = new Map<number, number>();
+	for (const id of validTokenIds) {
+		const e = Math.exp(rawLogits[id] - maxLogit);
+		expByToken.set(id, e);
+		sumExp += e;
+	}
+
+	// 2. Prefix expansion with max aggregation (probabilities sum to 1 over the
+	//    masked subset, so divide each token's exp by sumExp on the fly).
+	const wordProbs = new Map<string, number>();
+	for (const [id, words] of prefixMap) {
+		const p = sumExp > 0 ? (expByToken.get(id) ?? 0) / sumExp : 0;
+		for (const w of words) {
+			const prev = wordProbs.get(w) ?? 0;
+			if (p > prev) {
+				wordProbs.set(w, p);
+			}
+		}
+	}
+
+	// 3. Split into L2 / L3 and reserve the top L2 slots unconditionally.
+	const l2Matches: Array<[string, number]> = [];
+	const l3Matches: Array<[string, number]> = [];
+	for (const [w, p] of wordProbs) {
+		if (domainWords.has(w)) {
+			l2Matches.push([w, p]);
+		} else {
+			l3Matches.push([w, p]);
+		}
+	}
+	l2Matches.sort((a, b) => b[1] - a[1]);
+	const reserved = l2Matches.slice(0, BE_PARITY.RESERVED_L2_SLOTS);
+
+	// 4. Pool the leftovers in log space; the L2 bias only affects ranking here.
+	// Words in l2Matches are unique and the array is sorted descending, so the
+	// non-reserved entries are exactly the tail after the reserved prefix — slice
+	// it directly rather than allocating a Set and scanning every entry on this
+	// hot path (runs ~every word boundary while typing).
+	const pool: Array<[string, number]> = [];
+	for (const [w, p] of l2Matches.slice(BE_PARITY.RESERVED_L2_SLOTS)) {
+		pool.push([w, Math.log(Math.max(p, 1e-10)) + BE_PARITY.L2_BIAS]);
+	}
+	for (const [w, p] of l3Matches) {
+		pool.push([w, Math.log(Math.max(p, 1e-10))]);
+	}
+	pool.sort((a, b) => b[1] - a[1]);
+	const remainingSlots = Math.max(0, BE_PARITY.TOP_K_WORDS - reserved.length);
+	const poolWinners = pool.slice(0, remainingSlots);
+
+	// 5. Assemble payload: store RAW probabilities (the bias was ranking-only),
+	// lowercase keys, trimmed at MIN_PROB. Reserved first, then pool winners.
+	// Reserved entries are written first; pool-winner writes must NOT clobber a
+	// reserved entry whose normalised key collides (two source words can
+	// `.trim().toLowerCase()` to the same key — e.g. "Function" vs "function ").
+	// Without the existence guard, a low-probability pool winner would silently
+	// overwrite the (higher-probability) reserved entry, degrading top-K
+	// quality in a way that's invisible from the debug summary.
+	const result: Record<string, number> = {};
+	const addEntry = (word: string, prob: number, allowOverwrite: boolean): void => {
+		if (prob <= BE_PARITY.MIN_PROB) {
+			return;
+		}
+		const key = word.trim().toLowerCase();
+		if (!allowOverwrite && key in result) {
+			return;
+		}
+		result[key] = prob;
+	};
+	for (const [w, p] of reserved) {
+		addEntry(w, p, true);
+	}
+	for (const [w] of poolWinners) {
+		addEntry(w, wordProbs.get(w) ?? 0, false);
+	}
+
+	if (isAutocompleteDebugEnabled()) {
+		const topReserved = reserved
+			.slice(0, 5)
+			.map(([w, p]) => `${w}:${(p * 100).toFixed(2)}%`)
+			.join(', ');
+		const topPool = poolWinners
+			.slice(0, 5)
+			.map(([w]) => `${w}:${((wordProbs.get(w) ?? 0) * 100).toFixed(2)}%`)
+			.join(', ');
+		// eslint-disable-next-line no-console
+		console.log(
+			'%c[computeBePayload] %c%d valid tokens → %d words expanded | L2: %d / L3: %d | reserved: %d | pool winners: %d | final: %d words\n  maxLogit(masked): %s | sumExp: %s\n  top reserved L2: %s\n  top pool: %s',
+			'color: #9c27b0; font-weight: bold;',
+			'color: inherit;',
+			validTokenIds.length,
+			wordProbs.size,
+			l2Matches.length,
+			l3Matches.length,
+			reserved.length,
+			poolWinners.length,
+			Object.keys(result).length,
+			maxLogit.toFixed(3),
+			sumExp.toFixed(1),
+			topReserved || '(none)',
+			topPool || '(none)',
+		);
+	}
+
+	return result;
+};
 
 // ─── Factory ─────────────────────────────────────────────────────────────────
 
@@ -127,7 +553,7 @@ export const createLocalSlowLaneClient = (
 		debounceMs = DEFAULT_DEBOUNCE_MS,
 		onUpdate,
 		onStatus,
-		modelId = LOCAL_MLC_MODEL_ID,
+		modelId = LOCAL_MLC_CAUSAL_MODEL_ID,
 		customModelConfig,
 	} = config;
 
@@ -143,6 +569,9 @@ export const createLocalSlowLaneClient = (
 	let initFailed = false;
 	let engine: MLCEngine | null = null;
 	let engineInitPromise: Promise<void> | null = null;
+	// Captures raw next-token logits from the LM's single decode step. Registered
+	// with the engine below; `lmLogitsCapture.captured` is consumed in a later step.
+	const lmLogitsCapture = new CapturingLogitProcessor();
 
 	const unloadEngine = (engineToUnload: MLCEngine): void => {
 		engineToUnload.unload().catch((error: unknown) => {
@@ -178,20 +607,25 @@ export const createLocalSlowLaneClient = (
 			if (isAutocompleteDebugEnabled()) {
 				// eslint-disable-next-line no-console
 				console.log(
-					`%c[LocalSlowLane] %c🚀 Initialising MLC engine with model: ${modelId}`,
+					`%c[LocalSlowLane] %c🚀 Initialising MLC engine with models: ${modelId} (LM) + ${LOCAL_MLC_EMBEDDING_MODEL_ID} (embedder)`,
 					'color: #9c27b0; font-weight: bold;',
 					'color: inherit;',
 				);
 			}
-			onStatus?.(`Initialising model: ${modelId}…`);
+			onStatus?.(`Initialising models: ${modelId} + ${LOCAL_MLC_EMBEDDING_MODEL_ID}…`);
 
 			if (!('gpu' in navigator)) {
 				throw new Error('WebGPU not supported');
 			}
 
-			const { CreateMLCEngine, prebuiltAppConfig } = await import(
-				/* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-mlc-web-llm" */ '@mlc-ai/web-llm'
-			);
+			// Fetch the web-llm runtime and the BE-parity lookup tables in parallel;
+			// both are dynamically imported so they stay out of the main editor chunk.
+			const [{ MLCEngine: MLCEngineCtor, prebuiltAppConfig }] = await Promise.all([
+				import(
+					/* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-mlc-web-llm" */ '@mlc-ai/web-llm'
+				),
+				loadBePayloadData(),
+			]);
 
 			const customModelRecord: WebLlmModelRecord | undefined = customModelConfig
 				? {
@@ -222,27 +656,49 @@ export const createLocalSlowLaneClient = (
 				],
 			};
 
-			engine = await CreateMLCEngine(modelId, {
+			// Construct the engine with the logit-capture processor registered for
+			// the causal LM only (the embedder never decodes tokens), then load
+			// both the LM and the embedder into the same engine (multi-model).
+			const newEngine = new MLCEngineCtor({
 				appConfig,
 				initProgressCallback,
+				logitProcessorRegistry: new Map([[modelId, lmLogitsCapture]]),
 			});
 
+			await newEngine.reload([modelId, LOCAL_MLC_EMBEDDING_MODEL_ID]);
+
 			if (destroyed) {
 				// destroy() was called while we were loading — clean up
-				unloadEngine(engine);
-				engine = null;
+				unloadEngine(newEngine);
 				return;
 			}
 
+			engine = newEngine;
 			ready = true;
 
 			if (isAutocompleteDebugEnabled()) {
 				// eslint-disable-next-line no-console
 				console.log(
-					'%c[LocalSlowLane] %c✅ MLC engine loaded and ready',
+					'%c[LocalSlowLane] %c✅ Both models loaded and ready',
 					'color: #9c27b0; font-weight: bold;',
 					'color: #4caf50;',
 				);
+				// One-time identity summary so you can confirm which models are active
+				// without digging through the init-progress scroll.
+				// eslint-disable-next-line no-console
+				console.log(
+					'%c[LocalSlowLane] %c🧠 Causal LM    →',
+					'color: #9c27b0; font-weight: bold;',
+					'color: #2196f3; font-weight: bold;',
+					modelId,
+				);
+				// eslint-disable-next-line no-console
+				console.log(
+					'%c[LocalSlowLane] %c🔢 Embedder     →',
+					'color: #9c27b0; font-weight: bold;',
+					'color: #009688; font-weight: bold;',
+					LOCAL_MLC_EMBEDDING_MODEL_ID,
+				);
 			}
 			onStatus?.('Model loaded and ready.');
 		} catch (err) {
@@ -271,83 +727,112 @@ export const createLocalSlowLaneClient = (
 	// ── Inference ──────────────────────────────────────────────────────────
 
 	/**
-	 * Run a single forward pass to extract next-token logit probabilities.
+	 * Run a single forward pass to produce the BE-parity slow-lane outputs.
 	 *
-	 * We use the chat completions API with `max_tokens: 1` and `logprobs: true`
-	 * to get the model's next-token distribution without generating text.
-	 * This is the cheapest possible inference call — a single forward pass.
+	 * Two calls run in parallel on the shared engine:
+	 *   - `completions.create({ max_tokens: 1 })` runs the causal LM for exactly
+	 *     one decode step. We ignore the generated text; the LogitProcessor
+	 *     captures the raw next-token logits during that step, which we turn into
+	 *     a whole-word payload via `computeBePayload`.
+	 *   - `embeddings.create(...)` runs the Arctic embedder to produce the real
+	 *     384-d semantic vector (passage-encoded; see `wrapForArctic`).
 	 */
 	const runInference = async (text: string, requestId: number): Promise<void> => {
 		if (!engine || destroyed) {
 			return;
 		}
 
-		try {
-			// Use chat completion with logprobs to get next-token distribution
-			const response = await engine.chat.completions.create({
-				messages: [
-					{
-						role: 'user',
-						content: text,
-					},
-				],
-				max_tokens: 1,
-				logprobs: true,
-				top_logprobs: 5,
-				temperature: 0,
-			});
+		// Clear the capture buffer so we read only this pass's logits. The engine
+		// serialises per-model requests and updateContext is debounced, so the
+		// latest request's decode step is the last to populate `captured` before
+		// we read it below; stale requests bail on the latestRequestId guard.
+		lmLogitsCapture.resetState();
 
-			// Discard stale results
-			if (requestId < latestRequestId || destroyed) {
-				return;
-			}
+		// Apply BE-parity rolling-window truncation before both encoders.
+		// BE semantic: last max_context_words words (typeahead_context_encoding.py:36)
+		// BE causal LM: last max_context_tokens BPE tokens (causal_lm_encoder.py:194–198),
+		//               approximated here with word count (no tokenizer available on FE).
+		const lmText = truncateToLastNWords(text, BE_PARITY.MAX_CONTEXT_TOKENS);
+		const semanticText = truncateToLastNWords(text, BE_PARITY.MAX_CONTEXT_WORDS);
+		const arcticInput = wrapForArctic(semanticText);
 
-			// ── Extract LM logits ───────────────────────────────────────
-			const lmLogits: Record<string, number> = {};
+		if (isAutocompleteDebugEnabled()) {
+			// eslint-disable-next-line no-console
+			console.log(
+				`%c[LocalSlowLane] %c🔢 Arctic input (${arcticInput.length} chars, ${semanticText.split(/\s+/u).length} words): "${arcticInput.length > 100 ? `${arcticInput.slice(0, 100)}…` : arcticInput}"`,
+				'color: #9c27b0; font-weight: bold;',
+				'color: #009688;',
+			);
+			// eslint-disable-next-line no-console
+			console.log(
+				`%c[LocalSlowLane] %c🧠 LM input (${lmText.length} chars, ${lmText.split(/\s+/u).length} words): "${lmText.length > 100 ? `${lmText.slice(0, 100)}…` : lmText}"`,
+				'color: #9c27b0; font-weight: bold;',
+				'color: #2196f3;',
+			);
+		}
 
-			const logprobsContent = response.choices?.[0]?.logprobs?.content;
-			if (logprobsContent && logprobsContent.length > 0) {
-				const tokenLogprobs = logprobsContent[0];
+		try {
+			const tStart = performance.now();
+			let tLmDone = 0;
+			let tEmbDone = 0;
+
+			const [, embeddingResponse] = await Promise.all([
+				engine.completions
+					.create({
+						model: modelId,
+						prompt: lmText,
+						max_tokens: 1,
+						temperature: 0,
+						logprobs: false,
+					})
+					.then((r) => {
+						tLmDone = performance.now();
+						return r;
+					}),
+				engine.embeddings
+					.create({
+						model: LOCAL_MLC_EMBEDDING_MODEL_ID,
+						input: arcticInput,
+					})
+					.then((r) => {
+						tEmbDone = performance.now();
+						return r;
+					}),
+			]);
 
-				// Add the top token
-				if (tokenLogprobs.token) {
-					const token = tokenLogprobs.token.trim().toLowerCase();
-					// @ts-ignore TS1501: Unicode regex flag requires a newer TS target than the declaration build uses.
-					if (token.length > 0 && /^[a-z]/iu.test(token)) {
-						lmLogits[token] = Math.exp(tokenLogprobs.logprob);
-					}
-				}
+			if (isAutocompleteDebugEnabled()) {
+				// eslint-disable-next-line no-console
+				console.log(
+					`%c[LocalSlowLane] %c⏱ LM: ${(tLmDone - tStart).toFixed(0)}ms | Embedder: ${(tEmbDone - tStart).toFixed(0)}ms | Total: ${(Math.max(tLmDone, tEmbDone) - tStart).toFixed(0)}ms`,
+					'color: #9c27b0; font-weight: bold;',
+					'color: #ff9800;',
+				);
+			}
 
-				// Add alternative tokens from top_logprobs
-				if (tokenLogprobs.top_logprobs) {
-					for (const alt of tokenLogprobs.top_logprobs) {
-						const token = alt.token.trim().toLowerCase();
-						// @ts-ignore TS1501: Unicode regex flag requires a newer TS target than the declaration build uses.
-						if (token.length > 0 && /^[a-z]/iu.test(token)) {
-							lmLogits[token] = Math.exp(alt.logprob);
-						}
-					}
-				}
+			// Discard stale results
+			if (requestId < latestRequestId || destroyed) {
+				return;
 			}
 
-			storedLmLogits = Object.keys(lmLogits).length > 0 ? lmLogits : null;
-
-			// ── Semantic vector ─────────────────────────────────────────
-			// SmolLM is a generative model, not an embedding model, so we
-			// don't get a true semantic vector. We generate a lightweight
-			// pseudo-embedding from the logit distribution for compatibility
-			// with the existing scoring pipeline.
-			//
-			// For a production implementation, you would use a dedicated
-			// embedding model (e.g. via web-llm's embeddings API with an
-			// embedding-specific model).
-			if (storedLmLogits) {
-				const logitValues = Object.values(storedLmLogits);
-				storedContextVector = new Float32Array(logitValues);
+			// ── LM logits: whole-word BE-parity payload ──────────────────
+			const rawLogits = lmLogitsCapture.captured;
+			if (rawLogits) {
+				const payload = computeBePayload(rawLogits, firstTokenToWords, l2Words);
+				storedLmLogits = Object.keys(payload).length > 0 ? payload : null;
 			} else {
-				storedContextVector = null;
+				storedLmLogits = null;
 			}
 
+			// ── Semantic vector: real 384-d Arctic embedding ─────────────
+			// Guard against base64-encoded responses (encoding_format: 'base64' would
+			// yield a string, and new Float32Array(string) silently produces an empty
+			// array, corrupting downstream cosine-similarity scoring).
+			const embedding = embeddingResponse.data?.[0]?.embedding;
+			storedContextVector =
+				Array.isArray(embedding) && embedding.length > 0
+					? new Float32Array(embedding as number[])
+					: null;
+
 			if (isAutocompleteDebugEnabled()) {
 				// eslint-disable-next-line no-console
 				console.groupCollapsed(
@@ -355,16 +840,23 @@ export const createLocalSlowLaneClient = (
 					'color: #9c27b0; font-weight: bold;',
 					'color: inherit;',
 				);
-				// eslint-disable-next-line no-console
-				console.log(
-					storedContextVector
-						? `✅ pseudo-vector: ${storedContextVector.length} dims`
-						: '❌ No vector',
-				);
+				if (storedContextVector) {
+					let sumSq = 0;
+					for (let i = 0; i < storedContextVector.length; i++) {
+						sumSq += storedContextVector[i] * storedContextVector[i];
+					}
+					// eslint-disable-next-line no-console
+					console.log(
+						`✅ semantic vector: ${storedContextVector.length} dims (L2 norm ${Math.sqrt(sumSq).toFixed(3)})`,
+					);
+				} else {
+					// eslint-disable-next-line no-console
+					console.log('❌ No vector');
+				}
 				// eslint-disable-next-line no-console
 				console.log(
 					storedLmLogits
-						? `✅ lm_logits: ${Object.keys(storedLmLogits).length} tokens`
+						? `✅ lm_logits: ${Object.keys(storedLmLogits).length} words`
 						: '❌ No lm_logits',
 				);
 				if (storedLmLogits) {