@atlaskit/editor-plugin-autocomplete 3.1.0 → 3.2.0

package/dist/esm/pm-plugins/local-slow-lane-client.js

@@ -1,28 +1,40 @@
-import _slicedToArray from "@babel/runtime/helpers/slicedToArray";
 import _toConsumableArray from "@babel/runtime/helpers/toConsumableArray";
-import _defineProperty from "@babel/runtime/helpers/defineProperty";
+import _slicedToArray from "@babel/runtime/helpers/slicedToArray";
 import _asyncToGenerator from "@babel/runtime/helpers/asyncToGenerator";
+import _typeof from "@babel/runtime/helpers/typeof";
+import _createClass from "@babel/runtime/helpers/createClass";
+import _classCallCheck from "@babel/runtime/helpers/classCallCheck";
+import _defineProperty from "@babel/runtime/helpers/defineProperty";
+function ownKeys(e, r) { var t = Object.keys(e); if (Object.getOwnPropertySymbols) { var o = Object.getOwnPropertySymbols(e); r && (o = o.filter(function (r) { return Object.getOwnPropertyDescriptor(e, r).enumerable; })), t.push.apply(t, o); } return t; }
+function _objectSpread(e) { for (var r = 1; r < arguments.length; r++) { var t = null != arguments[r] ? arguments[r] : {}; r % 2 ? ownKeys(Object(t), !0).forEach(function (r) { _defineProperty(e, r, t[r]); }) : Object.getOwnPropertyDescriptors ? Object.defineProperties(e, Object.getOwnPropertyDescriptors(t)) : ownKeys(Object(t)).forEach(function (r) { Object.defineProperty(e, r, Object.getOwnPropertyDescriptor(t, r)); }); } return e; }
+import _regeneratorRuntime from "@babel/runtime/regenerator";
 function _createForOfIteratorHelper(r, e) { var t = "undefined" != typeof Symbol && r[Symbol.iterator] || r["@@iterator"]; if (!t) { if (Array.isArray(r) || (t = _unsupportedIterableToArray(r)) || e && r && "number" == typeof r.length) { t && (r = t); var _n = 0, F = function F() {}; return { s: F, n: function n() { return _n >= r.length ? { done: !0 } : { done: !1, value: r[_n++] }; }, e: function e(r) { throw r; }, f: F }; } throw new TypeError("Invalid attempt to iterate non-iterable instance.\nIn order to be iterable, non-array objects must have a [Symbol.iterator]() method."); } var o, a = !0, u = !1; return { s: function s() { t = t.call(r); }, n: function n() { var r = t.next(); return a = r.done, r; }, e: function e(r) { u = !0, o = r; }, f: function f() { try { a || null == t.return || t.return(); } finally { if (u) throw o; } } }; }
 function _unsupportedIterableToArray(r, a) { if (r) { if ("string" == typeof r) return _arrayLikeToArray(r, a); var t = {}.toString.call(r).slice(8, -1); return "Object" === t && r.constructor && (t = r.constructor.name), "Map" === t || "Set" === t ? Array.from(r) : "Arguments" === t || /^(?:Ui|I)nt(?:8|16|32)(?:Clamped)?Array$/.test(t) ? _arrayLikeToArray(r, a) : void 0; } }
 function _arrayLikeToArray(r, a) { (null == a || a > r.length) && (a = r.length); for (var e = 0, n = Array(a); e < a; e++) n[e] = r[e]; return n; }
-import _regeneratorRuntime from "@babel/runtime/regenerator";
-function ownKeys(e, r) { var t = Object.keys(e); if (Object.getOwnPropertySymbols) { var o = Object.getOwnPropertySymbols(e); r && (o = o.filter(function (r) { return Object.getOwnPropertyDescriptor(e, r).enumerable; })), t.push.apply(t, o); } return t; }
-function _objectSpread(e) { for (var r = 1; r < arguments.length; r++) { var t = null != arguments[r] ? arguments[r] : {}; r % 2 ? ownKeys(Object(t), !0).forEach(function (r) { _defineProperty(e, r, t[r]); }) : Object.getOwnPropertyDescriptors ? Object.defineProperties(e, Object.getOwnPropertyDescriptors(t)) : ownKeys(Object(t)).forEach(function (r) { Object.defineProperty(e, r, Object.getOwnPropertyDescriptor(t, r)); }); } return e; }
 /**
  * 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)
@@ -46,17 +58,498 @@ import { isWordBoundary } from './slow-lane-client';
 // ─── Constants ───────────────────────────────────────────────────────────────
 
 var DEFAULT_DEBOUNCE_MS = 300;
-export var LOCAL_MLC_MODEL_ID = 'SmolLM2-135M-Instruct-q0f16-MLC';
+export var 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 var 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 var wrapForArctic = function wrapForArctic(text) {
+  return "[CLS] ".concat(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 var 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
+};
+
+/**
+ * Return the last `n` whitespace-separated words of `text`, joined by spaces.
+ * Mirrors the BE rolling-window truncation applied before both encoders.
+ */
+var truncateToLastNWords = function truncateToLastNWords(text, n) {
+  var words = text.trim().split(/[\t-\r \xA0\u1680\u2000-\u200A\u2028\u2029\u202F\u205F\u3000\uFEFF]+/);
+  return words.length <= n ? text : words.slice(-n).join(' ');
+};
+
+// ─── Logit capture ─────────────────────────────────────────────────────────
+
+/**
+ * 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.
+ */
+var CapturingLogitProcessor = /*#__PURE__*/_createClass(function CapturingLogitProcessor() {
+  var _this = this;
+  _classCallCheck(this, CapturingLogitProcessor);
+  _defineProperty(this, "captured", null);
+  _defineProperty(this, "processLogits", function (logits) {
+    // Copy off web-llm's shared buffer — it may reuse `logits` across calls.
+    _this.captured = new Float32Array(logits);
+    return logits;
+  });
+  _defineProperty(this, "processSampledToken", function () {
+    // No-op — we don't track sampled tokens.
+  });
+  _defineProperty(this, "resetState", function () {
+    _this.captured = null;
+  });
+}); // ─── BE-parity data + algorithm ──────────────────────────────────────────────
+/**
+ * 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).
+ */
+var firstTokenToWords = 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()`.
+ */
+var l2Words = 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.
+ */
+var prefixMapTokenIds = [];
+
+/** De-dupes concurrent loads and lets repeated calls await the same payload. */
+var bePayloadDataPromise;
+
+/**
+ * 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.
+ */
+var unwrapJsonModule = function unwrapJsonModule(mod, shape) {
+  if (mod == null || _typeof(mod) !== 'object') {
+    return null;
+  }
+  var namespace = mod;
+
+  // Compute the named-export own-keys (strip synthetic markers).
+  var ownKeys = Object.keys(namespace).filter(function (k) {
+    return 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.
+      var len = ownKeys.length;
+      var arr = new Array(len);
+      for (var i = 0; i < len; i++) {
+        arr[i] = namespace[String(i)];
+      }
+      return arr;
+    }
+    // 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).
+    var obj = {};
+    var _iterator = _createForOfIteratorHelper(ownKeys),
+      _step;
+    try {
+      for (_iterator.s(); !(_step = _iterator.n()).done;) {
+        var k = _step.value;
+        obj[k] = namespace[k];
+      }
+    } catch (err) {
+      _iterator.e(err);
+    } finally {
+      _iterator.f();
+    }
+    return obj;
+  }
+
+  // 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;
+  }
+  return null;
+};
 
-/** HF root for the default weights (includes `tensor-cache.json` for WebLLM 0.2+). */
-export var LOCAL_MLC_HF_MODEL_REPO = 'https://huggingface.co/mlc-ai/SmolLM2-135M-Instruct-q0f16-MLC';
-export var LOCAL_MLC_MODEL_LIB_WASM_NAME = 'SmolLM2-135M-Instruct-q0f16-ctx4k_cs1k-webgpu.wasm';
+/**
+ * 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.
+ */
+var loadBePayloadData = function loadBePayloadData() {
+  if (!bePayloadDataPromise) {
+    bePayloadDataPromise = _asyncToGenerator( /*#__PURE__*/_regeneratorRuntime.mark(function _callee() {
+      var _yield$Promise$all, _yield$Promise$all2, firstTokenToWordsModule, vocabularyModule, firstTokenToWordsData, vocabularyData;
+      return _regeneratorRuntime.wrap(function (_context) {
+        while (1) switch (_context.prev = _context.next) {
+          case 0:
+            _context.next = 1;
+            return 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')]);
+          case 1:
+            _yield$Promise$all = _context.sent;
+            _yield$Promise$all2 = _slicedToArray(_yield$Promise$all, 2);
+            firstTokenToWordsModule = _yield$Promise$all2[0];
+            vocabularyModule = _yield$Promise$all2[1];
+            firstTokenToWordsData = unwrapJsonModule(firstTokenToWordsModule, 'object');
+            vocabularyData = unwrapJsonModule(vocabularyModule, 'object');
+            if (!(firstTokenToWordsData == null || (vocabularyData === null || vocabularyData === void 0 ? void 0 : vocabularyData.words) == null)) {
+              _context.next = 2;
+              break;
+            }
+            throw new Error("[LocalSlowLane] JSON module could not be unwrapped \u2014 " + "firstTokenToWordsData=".concat(firstTokenToWordsData == null ? 'null/undefined' : 'defined', ", ") + "vocabularyData=".concat(vocabularyData == null ? 'null/undefined' : vocabularyData.words == null ? 'defined but missing .words' : 'defined'));
+          case 2:
+            firstTokenToWords = new Map(Object.entries(firstTokenToWordsData).map(function (_ref2) {
+              var _ref3 = _slicedToArray(_ref2, 2),
+                tokenId = _ref3[0],
+                words = _ref3[1];
+              return [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
+              });
+            }
+          case 3:
+          case "end":
+            return _context.stop();
+        }
+      }, _callee);
+    }))().catch(function (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;
+};
 
 /**
- * Original target repo (add-basics fine-tune). **Not compatible with WebLLM 0.2.x** (no `tensor-cache.json`).
- * @see module doc above
+ * 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 var HUGGINGFACE_TB_SMOLLM_ADD_BASICS_REPO = 'https://huggingface.co/HuggingFaceTB/smollm-135M-instruct-add-basics-q0f16-MLC';
+export var computeBePayload = function computeBePayload(rawLogits, prefixMap, domainWords) {
+  var validTokenIds = arguments.length > 3 && arguments[3] !== undefined ? arguments[3] : prefixMapTokenIds;
+  // 1. Numerically-stable masked softmax over validTokenIds only.
+  var maxLogit = -Infinity;
+  var _iterator2 = _createForOfIteratorHelper(validTokenIds),
+    _step2;
+  try {
+    for (_iterator2.s(); !(_step2 = _iterator2.n()).done;) {
+      var id = _step2.value;
+      var v = rawLogits[id];
+      if (v > maxLogit) {
+        maxLogit = v;
+      }
+    }
+  } catch (err) {
+    _iterator2.e(err);
+  } finally {
+    _iterator2.f();
+  }
+  var sumExp = 0;
+  var expByToken = new Map();
+  var _iterator3 = _createForOfIteratorHelper(validTokenIds),
+    _step3;
+  try {
+    for (_iterator3.s(); !(_step3 = _iterator3.n()).done;) {
+      var _id = _step3.value;
+      var 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).
+  } catch (err) {
+    _iterator3.e(err);
+  } finally {
+    _iterator3.f();
+  }
+  var wordProbs = new Map();
+  var _iterator4 = _createForOfIteratorHelper(prefixMap),
+    _step4;
+  try {
+    for (_iterator4.s(); !(_step4 = _iterator4.n()).done;) {
+      var _expByToken$get;
+      var _step4$value = _slicedToArray(_step4.value, 2),
+        _id2 = _step4$value[0],
+        words = _step4$value[1];
+      var _p = sumExp > 0 ? ((_expByToken$get = expByToken.get(_id2)) !== null && _expByToken$get !== void 0 ? _expByToken$get : 0) / sumExp : 0;
+      var _iterator9 = _createForOfIteratorHelper(words),
+        _step9;
+      try {
+        for (_iterator9.s(); !(_step9 = _iterator9.n()).done;) {
+          var _wordProbs$get2;
+          var _w = _step9.value;
+          var prev = (_wordProbs$get2 = wordProbs.get(_w)) !== null && _wordProbs$get2 !== void 0 ? _wordProbs$get2 : 0;
+          if (_p > prev) {
+            wordProbs.set(_w, _p);
+          }
+        }
+      } catch (err) {
+        _iterator9.e(err);
+      } finally {
+        _iterator9.f();
+      }
+    }
+
+    // 3. Split into L2 / L3 and reserve the top L2 slots unconditionally.
+  } catch (err) {
+    _iterator4.e(err);
+  } finally {
+    _iterator4.f();
+  }
+  var l2Matches = [];
+  var l3Matches = [];
+  var _iterator5 = _createForOfIteratorHelper(wordProbs),
+    _step5;
+  try {
+    for (_iterator5.s(); !(_step5 = _iterator5.n()).done;) {
+      var _step5$value = _slicedToArray(_step5.value, 2),
+        _w2 = _step5$value[0],
+        _p2 = _step5$value[1];
+      if (domainWords.has(_w2)) {
+        l2Matches.push([_w2, _p2]);
+      } else {
+        l3Matches.push([_w2, _p2]);
+      }
+    }
+  } catch (err) {
+    _iterator5.e(err);
+  } finally {
+    _iterator5.f();
+  }
+  l2Matches.sort(function (a, b) {
+    return b[1] - a[1];
+  });
+  var 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).
+  var pool = [];
+  var _iterator6 = _createForOfIteratorHelper(l2Matches.slice(BE_PARITY.RESERVED_L2_SLOTS)),
+    _step6;
+  try {
+    for (_iterator6.s(); !(_step6 = _iterator6.n()).done;) {
+      var _step6$value = _slicedToArray(_step6.value, 2),
+        _w3 = _step6$value[0],
+        _p3 = _step6$value[1];
+      pool.push([_w3, Math.log(Math.max(_p3, 1e-10)) + BE_PARITY.L2_BIAS]);
+    }
+  } catch (err) {
+    _iterator6.e(err);
+  } finally {
+    _iterator6.f();
+  }
+  for (var _i = 0, _l3Matches = l3Matches; _i < _l3Matches.length; _i++) {
+    var _l3Matches$_i = _slicedToArray(_l3Matches[_i], 2),
+      w = _l3Matches$_i[0],
+      p = _l3Matches$_i[1];
+    pool.push([w, Math.log(Math.max(p, 1e-10))]);
+  }
+  pool.sort(function (a, b) {
+    return b[1] - a[1];
+  });
+  var remainingSlots = Math.max(0, BE_PARITY.TOP_K_WORDS - reserved.length);
+  var 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.
+  var result = {};
+  var addEntry = function addEntry(word, prob, allowOverwrite) {
+    if (prob <= BE_PARITY.MIN_PROB) {
+      return;
+    }
+    var key = word.trim().toLowerCase();
+    if (!allowOverwrite && key in result) {
+      return;
+    }
+    result[key] = prob;
+  };
+  var _iterator7 = _createForOfIteratorHelper(reserved),
+    _step7;
+  try {
+    for (_iterator7.s(); !(_step7 = _iterator7.n()).done;) {
+      var _step7$value = _slicedToArray(_step7.value, 2),
+        _w4 = _step7$value[0],
+        _p4 = _step7$value[1];
+      addEntry(_w4, _p4, true);
+    }
+  } catch (err) {
+    _iterator7.e(err);
+  } finally {
+    _iterator7.f();
+  }
+  var _iterator8 = _createForOfIteratorHelper(poolWinners),
+    _step8;
+  try {
+    for (_iterator8.s(); !(_step8 = _iterator8.n()).done;) {
+      var _wordProbs$get3;
+      var _step8$value = _slicedToArray(_step8.value, 1),
+        _w5 = _step8$value[0];
+      addEntry(_w5, (_wordProbs$get3 = wordProbs.get(_w5)) !== null && _wordProbs$get3 !== void 0 ? _wordProbs$get3 : 0, false);
+    }
+  } catch (err) {
+    _iterator8.e(err);
+  } finally {
+    _iterator8.f();
+  }
+  if (isAutocompleteDebugEnabled()) {
+    var topReserved = reserved.slice(0, 5).map(function (_ref4) {
+      var _ref5 = _slicedToArray(_ref4, 2),
+        w = _ref5[0],
+        p = _ref5[1];
+      return "".concat(w, ":").concat((p * 100).toFixed(2), "%");
+    }).join(', ');
+    var topPool = poolWinners.slice(0, 5).map(function (_ref6) {
+      var _wordProbs$get;
+      var _ref7 = _slicedToArray(_ref6, 1),
+        w = _ref7[0];
+      return "".concat(w, ":").concat((((_wordProbs$get = wordProbs.get(w)) !== null && _wordProbs$get !== void 0 ? _wordProbs$get : 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 ─────────────────────────────────────────────────────────────────
 
@@ -85,7 +578,7 @@ export var createLocalSlowLaneClient = function createLocalSlowLaneClient() {
     onUpdate = config.onUpdate,
     onStatus = config.onStatus,
     _config$modelId = config.modelId,
-    modelId = _config$modelId === void 0 ? LOCAL_MLC_MODEL_ID : _config$modelId,
+    modelId = _config$modelId === void 0 ? LOCAL_MLC_CAUSAL_MODEL_ID : _config$modelId,
     customModelConfig = config.customModelConfig;
 
   // ── State ──────────────────────────────────────────────────────────────
@@ -100,6 +593,9 @@ export var createLocalSlowLaneClient = function createLocalSlowLaneClient() {
   var initFailed = false;
   var engine = null;
   var engineInitPromise = 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.
+  var lmLogitsCapture = new CapturingLogitProcessor();
   var unloadEngine = function unloadEngine(engineToUnload) {
     engineToUnload.unload().catch(function (error) {
       if (isAutocompleteDebugEnabled()) {
@@ -120,29 +616,31 @@ export var createLocalSlowLaneClient = function createLocalSlowLaneClient() {
     onStatus === null || onStatus === void 0 || onStatus(message);
   };
   var initEngine = /*#__PURE__*/function () {
-    var _ref = _asyncToGenerator( /*#__PURE__*/_regeneratorRuntime.mark(function _callee() {
-      var _yield$import, CreateMLCEngine, prebuiltAppConfig, customModelRecord, appConfig, errorMsg, _t;
-      return _regeneratorRuntime.wrap(function (_context) {
-        while (1) switch (_context.prev = _context.next) {
+    var _ref8 = _asyncToGenerator( /*#__PURE__*/_regeneratorRuntime.mark(function _callee2() {
+      var _yield$Promise$all3, _yield$Promise$all4, _yield$Promise$all4$, MLCEngineCtor, prebuiltAppConfig, customModelRecord, appConfig, newEngine, errorMsg, _t;
+      return _regeneratorRuntime.wrap(function (_context2) {
+        while (1) switch (_context2.prev = _context2.next) {
           case 0:
-            _context.prev = 0;
+            _context2.prev = 0;
             if (isAutocompleteDebugEnabled()) {
               // eslint-disable-next-line no-console
-              console.log("%c[LocalSlowLane] %c\uD83D\uDE80 Initialising MLC engine with model: ".concat(modelId), 'color: #9c27b0; font-weight: bold;', 'color: inherit;');
+              console.log("%c[LocalSlowLane] %c\uD83D\uDE80 Initialising MLC engine with models: ".concat(modelId, " (LM) + ").concat(LOCAL_MLC_EMBEDDING_MODEL_ID, " (embedder)"), 'color: #9c27b0; font-weight: bold;', 'color: inherit;');
             }
-            onStatus === null || onStatus === void 0 || onStatus("Initialising model: ".concat(modelId, "\u2026"));
+            onStatus === null || onStatus === void 0 || onStatus("Initialising models: ".concat(modelId, " + ").concat(LOCAL_MLC_EMBEDDING_MODEL_ID, "\u2026"));
             if ('gpu' in navigator) {
-              _context.next = 1;
+              _context2.next = 1;
               break;
             }
             throw new Error('WebGPU not supported');
           case 1:
-            _context.next = 2;
-            return import( /* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-mlc-web-llm" */'@mlc-ai/web-llm');
+            _context2.next = 2;
+            return Promise.all([import( /* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-mlc-web-llm" */'@mlc-ai/web-llm'), loadBePayloadData()]);
           case 2:
-            _yield$import = _context.sent;
-            CreateMLCEngine = _yield$import.CreateMLCEngine;
-            prebuiltAppConfig = _yield$import.prebuiltAppConfig;
+            _yield$Promise$all3 = _context2.sent;
+            _yield$Promise$all4 = _slicedToArray(_yield$Promise$all3, 1);
+            _yield$Promise$all4$ = _yield$Promise$all4[0];
+            MLCEngineCtor = _yield$Promise$all4$.MLCEngine;
+            prebuiltAppConfig = _yield$Promise$all4$.prebuiltAppConfig;
             customModelRecord = customModelConfig ? _objectSpread(_objectSpread({
               model: customModelConfig.model,
               model_id: modelId,
@@ -158,34 +656,43 @@ export var createLocalSlowLaneClient = function createLocalSlowLaneClient() {
             } : {}) : undefined;
             appConfig = {
               model_list: [].concat(_toConsumableArray(prebuiltAppConfig.model_list), _toConsumableArray(customModelRecord ? [customModelRecord] : []))
-            };
-            _context.next = 3;
-            return 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).
+            newEngine = new MLCEngineCtor({
               appConfig: appConfig,
-              initProgressCallback: initProgressCallback
+              initProgressCallback: initProgressCallback,
+              logitProcessorRegistry: new Map([[modelId, lmLogitsCapture]])
             });
+            _context2.next = 3;
+            return newEngine.reload([modelId, LOCAL_MLC_EMBEDDING_MODEL_ID]);
           case 3:
-            engine = _context.sent;
             if (!destroyed) {
-              _context.next = 4;
+              _context2.next = 4;
               break;
             }
             // destroy() was called while we were loading — clean up
-            unloadEngine(engine);
-            engine = null;
-            return _context.abrupt("return");
+            unloadEngine(newEngine);
+            return _context2.abrupt("return");
           case 4:
+            engine = newEngine;
             ready = true;
             if (isAutocompleteDebugEnabled()) {
               // eslint-disable-next-line no-console
-              console.log('%c[LocalSlowLane] %c✅ MLC engine loaded and ready', 'color: #9c27b0; font-weight: bold;', 'color: #4caf50;');
+              console.log('%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 === null || onStatus === void 0 || onStatus('Model loaded and ready.');
-            _context.next = 6;
+            _context2.next = 6;
             break;
           case 5:
-            _context.prev = 5;
-            _t = _context["catch"](0);
+            _context2.prev = 5;
+            _t = _context2["catch"](0);
             errorMsg = _t instanceof Error ? _t.message : String(_t);
             ready = false;
             if (isAutocompleteDebugEnabled()) {
@@ -197,12 +704,12 @@ export var createLocalSlowLaneClient = function createLocalSlowLaneClient() {
             initFailed = true;
           case 6:
           case "end":
-            return _context.stop();
+            return _context2.stop();
         }
-      }, _callee, null, [[0, 5]]);
+      }, _callee2, null, [[0, 5]]);
     }));
     return function initEngine() {
-      return _ref.apply(this, arguments);
+      return _ref8.apply(this, arguments);
     };
   }();
   var ensureEngineInitialized = function ensureEngineInitialized() {
@@ -218,110 +725,127 @@ export var createLocalSlowLaneClient = function 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`).
    */
   var runInference = /*#__PURE__*/function () {
-    var _ref2 = _asyncToGenerator( /*#__PURE__*/_regeneratorRuntime.mark(function _callee2(text, requestId) {
-      var _response$choices, response, lmLogits, logprobsContent, tokenLogprobs, token, _iterator, _step, alt, _token, logitValues, topTokens, errorMsg, _t2;
-      return _regeneratorRuntime.wrap(function (_context2) {
-        while (1) switch (_context2.prev = _context2.next) {
+    var _ref9 = _asyncToGenerator( /*#__PURE__*/_regeneratorRuntime.mark(function _callee3(text, requestId) {
+      var lmText, semanticText, arcticInput, _embeddingResponse$da, tStart, tLmDone, tEmbDone, _yield$Promise$all5, _yield$Promise$all6, embeddingResponse, rawLogits, payload, embedding, sumSq, i, topTokens, errorMsg, _t2;
+      return _regeneratorRuntime.wrap(function (_context3) {
+        while (1) switch (_context3.prev = _context3.next) {
           case 0:
             if (!(!engine || destroyed)) {
-              _context2.next = 1;
+              _context3.next = 1;
               break;
             }
-            return _context2.abrupt("return");
+            return _context3.abrupt("return");
           case 1:
-            _context2.prev = 1;
-            _context2.next = 2;
-            return engine.chat.completions.create({
-              messages: [{
-                role: 'user',
-                content: text
-              }],
-              max_tokens: 1,
-              logprobs: true,
-              top_logprobs: 5,
-              temperature: 0
-            });
-          case 2:
-            response = _context2.sent;
-            if (!(requestId < latestRequestId || destroyed)) {
-              _context2.next = 3;
-              break;
+            // 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();
+
+            // 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).
+            lmText = truncateToLastNWords(text, BE_PARITY.MAX_CONTEXT_TOKENS);
+            semanticText = truncateToLastNWords(text, BE_PARITY.MAX_CONTEXT_WORDS);
+            arcticInput = wrapForArctic(semanticText);
+            if (isAutocompleteDebugEnabled()) {
+              // eslint-disable-next-line no-console
+              console.log("%c[LocalSlowLane] %c\uD83D\uDD22 Arctic input (".concat(arcticInput.length, " chars, ").concat(semanticText.split(/[\t-\r \xA0\u1680\u2000-\u200A\u2028\u2029\u202F\u205F\u3000\uFEFF]+/).length, " words): \"").concat(arcticInput.length > 100 ? "".concat(arcticInput.slice(0, 100), "\u2026") : arcticInput, "\""), 'color: #9c27b0; font-weight: bold;', 'color: #009688;');
+              // eslint-disable-next-line no-console
+              console.log("%c[LocalSlowLane] %c\uD83E\uDDE0 LM input (".concat(lmText.length, " chars, ").concat(lmText.split(/[\t-\r \xA0\u1680\u2000-\u200A\u2028\u2029\u202F\u205F\u3000\uFEFF]+/).length, " words): \"").concat(lmText.length > 100 ? "".concat(lmText.slice(0, 100), "\u2026") : lmText, "\""), 'color: #9c27b0; font-weight: bold;', 'color: #2196f3;');
             }
-            return _context2.abrupt("return");
+            _context3.prev = 2;
+            tStart = performance.now();
+            tLmDone = 0;
+            tEmbDone = 0;
+            _context3.next = 3;
+            return Promise.all([engine.completions.create({
+              model: modelId,
+              prompt: lmText,
+              max_tokens: 1,
+              temperature: 0,
+              logprobs: false
+            }).then(function (r) {
+              tLmDone = performance.now();
+              return r;
+            }), engine.embeddings.create({
+              model: LOCAL_MLC_EMBEDDING_MODEL_ID,
+              input: arcticInput
+            }).then(function (r) {
+              tEmbDone = performance.now();
+              return r;
+            })]);
           case 3:
-            // ── Extract LM logits ───────────────────────────────────────
-            lmLogits = {};
-            logprobsContent = (_response$choices = response.choices) === null || _response$choices === void 0 || (_response$choices = _response$choices[0]) === null || _response$choices === void 0 || (_response$choices = _response$choices.logprobs) === null || _response$choices === void 0 ? void 0 : _response$choices.content;
-            if (logprobsContent && logprobsContent.length > 0) {
-              tokenLogprobs = logprobsContent[0]; // Add the top token
-              if (tokenLogprobs.token) {
-                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\u017F\u212A]/i.test(token)) {
-                  lmLogits[token] = Math.exp(tokenLogprobs.logprob);
-                }
-              }
+            _yield$Promise$all5 = _context3.sent;
+            _yield$Promise$all6 = _slicedToArray(_yield$Promise$all5, 2);
+            embeddingResponse = _yield$Promise$all6[1];
+            if (isAutocompleteDebugEnabled()) {
+              // eslint-disable-next-line no-console
+              console.log("%c[LocalSlowLane] %c\u23F1 LM: ".concat((tLmDone - tStart).toFixed(0), "ms | Embedder: ").concat((tEmbDone - tStart).toFixed(0), "ms | Total: ").concat((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) {
-                _iterator = _createForOfIteratorHelper(tokenLogprobs.top_logprobs);
-                try {
-                  for (_iterator.s(); !(_step = _iterator.n()).done;) {
-                    alt = _step.value;
-                    _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\u017F\u212A]/i.test(_token)) {
-                      lmLogits[_token] = Math.exp(alt.logprob);
-                    }
-                  }
-                } catch (err) {
-                  _iterator.e(err);
-                } finally {
-                  _iterator.f();
-                }
-              }
+            // Discard stale results
+            if (!(requestId < latestRequestId || destroyed)) {
+              _context3.next = 4;
+              break;
             }
-            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) {
-              logitValues = Object.values(storedLmLogits);
-              storedContextVector = new Float32Array(logitValues);
+            return _context3.abrupt("return");
+          case 4:
+            // ── LM logits: whole-word BE-parity payload ──────────────────
+            rawLogits = lmLogitsCapture.captured;
+            if (rawLogits) {
+              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).
+            embedding = (_embeddingResponse$da = embeddingResponse.data) === null || _embeddingResponse$da === void 0 || (_embeddingResponse$da = _embeddingResponse$da[0]) === null || _embeddingResponse$da === void 0 ? void 0 : _embeddingResponse$da.embedding;
+            storedContextVector = Array.isArray(embedding) && embedding.length > 0 ? new Float32Array(embedding) : null;
             if (isAutocompleteDebugEnabled()) {
               // eslint-disable-next-line no-console
               console.groupCollapsed("%c[LocalSlowLane] %c\uD83D\uDCE5 Inference result (request #".concat(requestId, ")"), 'color: #9c27b0; font-weight: bold;', 'color: inherit;');
+              if (storedContextVector) {
+                sumSq = 0;
+                for (i = 0; i < storedContextVector.length; i++) {
+                  sumSq += storedContextVector[i] * storedContextVector[i];
+                }
+                // eslint-disable-next-line no-console
+                console.log("\u2705 semantic vector: ".concat(storedContextVector.length, " dims (L2 norm ").concat(Math.sqrt(sumSq).toFixed(3), ")"));
+              } else {
+                // eslint-disable-next-line no-console
+                console.log('❌ No vector');
+              }
               // eslint-disable-next-line no-console
-              console.log(storedContextVector ? "\u2705 pseudo-vector: ".concat(storedContextVector.length, " dims") : '❌ No vector');
-              // eslint-disable-next-line no-console
-              console.log(storedLmLogits ? "\u2705 lm_logits: ".concat(Object.keys(storedLmLogits).length, " tokens") : '❌ No lm_logits');
+              console.log(storedLmLogits ? "\u2705 lm_logits: ".concat(Object.keys(storedLmLogits).length, " words") : '❌ No lm_logits');
               if (storedLmLogits) {
-                topTokens = Object.entries(storedLmLogits).sort(function (_ref3, _ref4) {
-                  var _ref5 = _slicedToArray(_ref3, 2),
-                    a = _ref5[1];
-                  var _ref6 = _slicedToArray(_ref4, 2),
-                    b = _ref6[1];
+                topTokens = Object.entries(storedLmLogits).sort(function (_ref0, _ref1) {
+                  var _ref10 = _slicedToArray(_ref0, 2),
+                    a = _ref10[1];
+                  var _ref11 = _slicedToArray(_ref1, 2),
+                    b = _ref11[1];
                   return b - a;
                 }).slice(0, 10); // eslint-disable-next-line no-console
-                console.log('Top 10 predictions:', topTokens.map(function (_ref7) {
-                  var _ref8 = _slicedToArray(_ref7, 2),
-                    t = _ref8[0],
-                    p = _ref8[1];
+                console.log('Top 10 predictions:', topTokens.map(function (_ref12) {
+                  var _ref13 = _slicedToArray(_ref12, 2),
+                    t = _ref13[0],
+                    p = _ref13[1];
                   return "".concat(t, ": ").concat((p * 100).toFixed(1), "%");
                 }).join(', '));
               }
@@ -333,17 +857,17 @@ export var createLocalSlowLaneClient = function createLocalSlowLaneClient() {
               hasVector: storedContextVector !== null,
               hasLmLogits: storedLmLogits !== null
             });
-            _context2.next = 6;
+            _context3.next = 7;
             break;
-          case 4:
-            _context2.prev = 4;
-            _t2 = _context2["catch"](1);
+          case 5:
+            _context3.prev = 5;
+            _t2 = _context3["catch"](2);
             if (!(requestId < latestRequestId)) {
-              _context2.next = 5;
+              _context3.next = 6;
               break;
             }
-            return _context2.abrupt("return");
-          case 5:
+            return _context3.abrupt("return");
+          case 6:
             storedContextVector = null;
             storedLmLogits = null;
             onUpdate === null || onUpdate === void 0 || onUpdate({
@@ -356,14 +880,14 @@ export var createLocalSlowLaneClient = function createLocalSlowLaneClient() {
               // eslint-disable-next-line no-console
               console.log("%c[LocalSlowLane] %c\u274C Inference error (request #".concat(requestId, "): ").concat(errorMsg), 'color: #9c27b0; font-weight: bold;', 'color: #f44336;');
             }
-          case 6:
+          case 7:
           case "end":
-            return _context2.stop();
+            return _context3.stop();
         }
-      }, _callee2, null, [[1, 4]]);
+      }, _callee3, null, [[2, 5]]);
     }));
     return function runInference(_x, _x2) {
-      return _ref2.apply(this, arguments);
+      return _ref9.apply(this, arguments);
     };
   }();