@atlaskit/editor-plugin-autocomplete 3.0.0 → 3.2.0

package/dist/esm/pm-plugins/scoring-pipeline.js

@@ -161,20 +161,21 @@ function applyGrammarFilter(candidates, previousWord) {
         dropped.push(entry.candidate.word);
       }
     }
+
+    // Grammar is authoritative.
   } catch (err) {
     _iterator2.e(err);
   } finally {
     _iterator2.f();
   }
-  var finalFiltered = filtered.length > 0 ? filtered : candidates;
   return {
-    filtered: finalFiltered,
+    filtered: filtered,
     grammarMeta: {
       prevWord: lowerPrev,
       prevTags: prevTags,
       before: candidates.length,
-      after: finalFiltered.length,
-      dropped: filtered.length > 0 ? dropped : []
+      after: filtered.length,
+      dropped: dropped
     }
   };
 }

package/dist/esm/pm-plugins/text-predictor.js

@@ -1,4 +1,5 @@
 import _asyncToGenerator from "@babel/runtime/helpers/asyncToGenerator";
+import _typeof from "@babel/runtime/helpers/typeof";
 import _slicedToArray from "@babel/runtime/helpers/slicedToArray";
 import _createClass from "@babel/runtime/helpers/createClass";
 import _classCallCheck from "@babel/runtime/helpers/classCallCheck";
@@ -26,11 +27,9 @@ function _arrayLikeToArray(r, a) { (null == a || a > r.length) && (a = r.length)
 
 import { EXPERIENCE_NAME, failExp, startExp, succeedExp } from '../analytics/ufo';
 
-// import bigramsData from './data/bigrams.json';
-import l3VocabularyData from './data/l3_vocabulary.json';
-import vocabularyData from './data/vocabulary_10k.json';
-import wordIndexData from './data/word_index_10k.json';
-// import { rankCandidates, isGrammarAllowed } from './scoring-pipeline';
+// The vocabulary, L3 and word-index JSON payloads are dynamically imported in
+// loadDefaultVocabulary / loadVectorsAsync so their (large) contents stay out of
+// the editor's main chunk and only load when autocomplete is initialised.
 import { isAutocompleteDebugEnabled } from './debug-mode';
 import { rankCandidates, STAGE1_WEIGHT, STAGE2_WEIGHT, MIN_STAGE1_SCORE } from './scoring-pipeline';
 import { getStoredContextVector, getStoredLmLogits } from './slow-lane-client';
@@ -42,7 +41,7 @@ var PUNCTUATION_BOUNDARY_REGEX = /^[.,;:!?()\[\]{}"'`]+|[.,;:!?()\[\]{}"'`]+$/g;
 var MIN_PREFIX_LENGTH = 3;
 var MAX_CANDIDATES = 200;
 var CONTEXT_WORDS = 10;
-var MIN_SCORE_THRESHOLD = 0.2;
+var MIN_SCORE_THRESHOLD = 0.35;
 var L3_BASELINE_FREQ = 0.001;
 
 // ─── Types ───────────────────────────────────────────────────────────────────
@@ -206,7 +205,6 @@ var wordTrie = new WeightedWordTrie();
 // L3 Trie (General English Fallback)
 var l3Trie = new WeightedWordTrie();
 
-// --- Initialization Function ---
 /**
  * Loads the General English vocabulary.
  * expects a simple array of strings: ["about", "above", "actually", ...]
@@ -330,14 +328,11 @@ var tokenize = function tokenize(text) {
   return tokens;
 };
 var extractPreviousWord = function extractPreviousWord(text) {
-  // 1. Split the text by newlines or punctuation (. ? !)
+  // Only consider the current sentence/line the user is typing in.
   // eslint-disable-next-line require-unicode-regexp
   var sentences = text.split(/[\n.?!]+/);
-
-  // 2. Only look at the current sentence/line the user is typing in
   var currentSentence = sentences[sentences.length - 1];
 
-  // 3. Extract the previous word as normal
   // eslint-disable-next-line require-unicode-regexp
   var words = currentSentence.trimEnd().split(/\s+/);
   return words.length >= 2 ? words[words.length - 2] : '';
@@ -400,7 +395,6 @@ export var incrementSessionFreq = function incrementSessionFreq(word) {
  * Pass `undefined` (or omit the argument) to skip priming — useful when the
  * calling context does not yet have a page value available.
  */
-// NOTE: We ingest full page context here
 export var ingestDocumentPage = function ingestDocumentPage(pageContent) {
   if (!pageContent) {
     return;
@@ -433,7 +427,11 @@ export var ingestDocumentPage = function ingestDocumentPage(pageContent) {
 };
 export var predict = function predict(textBefore) {
   if (!isInitialized) {
-    loadDefaultVocabulary();
+    // Vocabulary JSON is code-split and loads asynchronously. Kick off the load
+    // and skip this keystroke; the plugin also primes it on focus, so the tries
+    // are usually ready before the user types.
+    void loadDefaultVocabulary().catch(function () {});
+    return null;
   }
   var t0 = performance.now();
 
@@ -485,14 +483,11 @@ export var predict = function predict(textBefore) {
   if (currentWord.length < MIN_PREFIX_LENGTH) {
     return null;
   }
-
-  // 1. Primary Query: Ask the L2 Domain Trie
   var candidates = wordTrie.getCandidates(currentWord, MAX_CANDIDATES);
 
-  // 2. Fallback Query: Gap-fill with the L3 General English Trie
+  // Gap-fill from the L3 general-English trie, requesting a full buffer so
+  // enough survive de-duplication against the L2 results.
   if (candidates.length < MAX_CANDIDATES) {
-    // Ask L3 for MAX_CANDIDATES to guarantee we have enough buffer
-    // to survive the deduplication process.
     var l3Candidates = l3Trie.getCandidates(currentWord, MAX_CANDIDATES);
     var existingWords = new Set(candidates.map(function (c) {
       return c.word;
@@ -502,8 +497,7 @@ export var predict = function predict(textBefore) {
     try {
       for (_iterator0.s(); !(_step0 = _iterator0.n()).done;) {
         var l3c = _step0.value;
-        if (candidates.length >= MAX_CANDIDATES) break; // Stop exactly at the limit
-
+        if (candidates.length >= MAX_CANDIDATES) break;
         if (!existingWords.has(l3c.word)) {
           candidates.push(l3c);
         }
@@ -687,9 +681,55 @@ export var predict = function predict(textBefore) {
 
 // ─── Data Loading ────────────────────────────────────────────────────────────
 
+/**
+ * Unwrap a dynamically imported JSON module to its parsed value, handling both
+ * interop modes AFM's bundler chain emits: a `.default`-wrapped namespace
+ * (classic webpack) and a named-exports namespace (webpack 5 / atlaspack JSON
+ * modules, where `default` can be a misleading scalar). Named exports are
+ * preferred when present. The caller declares the JSON `shape` because a dense
+ * array and a sparse numeric-keyed object are emitted identically as named
+ * exports. Kept in lock-step with the matching helper in local-slow-lane-client.ts.
+ */
+var unwrapJsonModule = function unwrapJsonModule(mod, shape) {
+  if (mod == null || _typeof(mod) !== 'object') {
+    return null;
+  }
+  var namespace = mod;
+  var ownKeys = Object.keys(namespace).filter(function (k) {
+    return k !== 'default' && k !== '__esModule';
+  });
+  if (ownKeys.length > 0) {
+    if (shape === 'array') {
+      var len = ownKeys.length;
+      var arr = new Array(len);
+      for (var i = 0; i < len; i++) {
+        arr[i] = namespace[String(i)];
+      }
+      return arr;
+    }
+    var obj = {};
+    var _iterator1 = _createForOfIteratorHelper(ownKeys),
+      _step1;
+    try {
+      for (_iterator1.s(); !(_step1 = _iterator1.n()).done;) {
+        var k = _step1.value;
+        obj[k] = namespace[k];
+      }
+    } catch (err) {
+      _iterator1.e(err);
+    } finally {
+      _iterator1.f();
+    }
+    return obj;
+  }
+  if ('default' in namespace && namespace.default != null) {
+    return namespace.default;
+  }
+  return null;
+};
 export var loadVectorsAsync = /*#__PURE__*/function () {
   var _ref6 = _asyncToGenerator( /*#__PURE__*/_regeneratorRuntime.mark(function _callee(options) {
-    var url, res, buffer, float32, wordIndex, nWords, dim, _t, _t2;
+    var url, _wordIndexOuter$index, res, buffer, float32, wordIndexModule, wordIndexOuter, wordIndex, nWords, dim, _t, _t2;
     return _regeneratorRuntime.wrap(function (_context) {
       while (1) switch (_context.prev = _context.next) {
         case 0:
@@ -749,9 +789,20 @@ export var loadVectorsAsync = /*#__PURE__*/function () {
           return res.arrayBuffer();
         case 9:
           buffer = _context.sent;
-          float32 = new Float32Array(buffer);
-          wordIndex = wordIndexData;
+          float32 = new Float32Array(buffer); // word_index_10k.json is wrapped as `{ "index": {…} }` so no real entry
+          // (e.g. the word "default") can shadow the synthetic ESM `default` export
+          // the bundler creates for dynamically-imported JSON.
+          _context.next = 10;
+          return import( /* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-word-index-10k" */'./data/word_index_10k.json');
+        case 10:
+          wordIndexModule = _context.sent;
+          wordIndexOuter = unwrapJsonModule(wordIndexModule, 'object');
+          wordIndex = (_wordIndexOuter$index = wordIndexOuter === null || wordIndexOuter === void 0 ? void 0 : wordIndexOuter.index) !== null && _wordIndexOuter$index !== void 0 ? _wordIndexOuter$index : {};
           nWords = Object.keys(wordIndex).length;
+          if (nWords === 0) {
+            // eslint-disable-next-line no-console
+            console.warn('[text-predictor] word_index_10k.json missing its `index` wrapper — wordIndex is empty, semantic scoring will be a no-op.');
+          }
           dim = float32.length / nWords;
           vectorStore = {
             float32: float32,
@@ -771,10 +822,10 @@ export var loadVectorsAsync = /*#__PURE__*/function () {
               sizeBytes: float32.byteLength
             });
           }
-          _context.next = 11;
+          _context.next = 12;
           break;
-        case 10:
-          _context.prev = 10;
+        case 11:
+          _context.prev = 11;
           _t2 = _context["catch"](6);
           vectorsLoadStarted = false;
           failExp(EXPERIENCE_NAME.LOAD_VECTORS, 'singleton', {
@@ -782,11 +833,11 @@ export var loadVectorsAsync = /*#__PURE__*/function () {
           });
           // eslint-disable-next-line no-console
           console.warn('[text-predictor] Failed to load vectors:', _t2);
-        case 11:
+        case 12:
         case "end":
           return _context.stop();
       }
-    }, _callee, null, [[3, 5], [6, 10]]);
+    }, _callee, null, [[3, 5], [6, 11]]);
   }));
   return function loadVectorsAsync(_x) {
     return _ref6.apply(this, arguments);
@@ -795,40 +846,73 @@ export var loadVectorsAsync = /*#__PURE__*/function () {
 export var initVectors = function initVectors(store) {
   vectorStore = store;
 };
+var vocabularyLoadPromise;
 export var loadDefaultVocabulary = function loadDefaultVocabulary() {
   if (isInitialized) {
-    return;
+    return Promise.resolve();
   }
-  startExp(EXPERIENCE_NAME.LOAD_VOCABULARY, 'singleton');
-  try {
-    // 1. Load the Atlassian Domain (L2)
-    var data = vocabularyData;
-    var terms = Object.entries(data.words).map(function (_ref7) {
-      var _ref8 = _slicedToArray(_ref7, 2),
-        word = _ref8[0],
-        stats = _ref8[1];
-      return {
-        word: word,
-        freq: stats.freq,
-        docFreq: stats.doc_freq,
-        authorFreq: stats.author_freq
-      };
-    });
-    initVocabulary({
-      terms: terms
-    });
-
-    // 2. Load General English (L3)
-    var l3Words = l3VocabularyData;
-    initL3Vocabulary(l3Words);
-    succeedExp(EXPERIENCE_NAME.LOAD_VOCABULARY, 'singleton', {
-      l2WordCount: terms.length,
-      l3WordCount: l3Words.length
-    });
-  } catch (e) {
-    failExp(EXPERIENCE_NAME.LOAD_VOCABULARY, 'singleton', {
-      errorType: 'parse_error'
-    });
-    throw e;
+  if (vocabularyLoadPromise) {
+    return vocabularyLoadPromise;
   }
+  vocabularyLoadPromise = _asyncToGenerator( /*#__PURE__*/_regeneratorRuntime.mark(function _callee2() {
+    var _yield$Promise$all, _yield$Promise$all2, vocabularyModule, l3VocabularyModule, vocabularyData, l3VocabularyData, terms, _t3;
+    return _regeneratorRuntime.wrap(function (_context2) {
+      while (1) switch (_context2.prev = _context2.next) {
+        case 0:
+          startExp(EXPERIENCE_NAME.LOAD_VOCABULARY, 'singleton');
+          _context2.prev = 1;
+          _context2.next = 2;
+          return Promise.all([import( /* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-vocabulary-10k" */'./data/vocabulary_10k.json'), import( /* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-l3-vocabulary" */'./data/l3_vocabulary.json')]);
+        case 2:
+          _yield$Promise$all = _context2.sent;
+          _yield$Promise$all2 = _slicedToArray(_yield$Promise$all, 2);
+          vocabularyModule = _yield$Promise$all2[0];
+          l3VocabularyModule = _yield$Promise$all2[1];
+          vocabularyData = unwrapJsonModule(vocabularyModule, 'object');
+          l3VocabularyData = unwrapJsonModule(l3VocabularyModule, 'array');
+          if (!((vocabularyData === null || vocabularyData === void 0 ? void 0 : vocabularyData.words) == null || !Array.isArray(l3VocabularyData))) {
+            _context2.next = 3;
+            break;
+          }
+          throw new Error('[text-predictor] vocabulary JSON modules could not be unwrapped');
+        case 3:
+          terms = Object.entries(vocabularyData.words).map(function (_ref8) {
+            var _ref9 = _slicedToArray(_ref8, 2),
+              word = _ref9[0],
+              stats = _ref9[1];
+            return {
+              word: word,
+              freq: stats.freq,
+              docFreq: stats.doc_freq,
+              authorFreq: stats.author_freq
+            };
+          }); // Load L3 before L2: initVocabulary flips `isInitialized = true`, so it
+          // must run last — otherwise a throw in initL3Vocabulary would strand
+          // `isInitialized` true and the retry path could never reload L3.
+          initL3Vocabulary(l3VocabularyData);
+          initVocabulary({
+            terms: terms
+          });
+          succeedExp(EXPERIENCE_NAME.LOAD_VOCABULARY, 'singleton', {
+            l2WordCount: terms.length,
+            l3WordCount: l3VocabularyData.length
+          });
+          _context2.next = 5;
+          break;
+        case 4:
+          _context2.prev = 4;
+          _t3 = _context2["catch"](1);
+          failExp(EXPERIENCE_NAME.LOAD_VOCABULARY, 'singleton', {
+            errorType: 'parse_error'
+          });
+          // Allow a later call to retry the load rather than caching the failure.
+          vocabularyLoadPromise = undefined;
+          throw _t3;
+        case 5:
+        case "end":
+          return _context2.stop();
+      }
+    }, _callee2, null, [[1, 4]]);
+  }))();
+  return vocabularyLoadPromise;
 };

package/dist/types/pm-plugins/local-slow-lane-client.d.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)
@@ -72,15 +81,93 @@ export interface LocalSlowLaneClient {
     setLmLogits: (logits: Record<string, number> | null) => void;
     updateContext: (text: string) => void;
 }
-export declare const LOCAL_MLC_MODEL_ID = "SmolLM2-135M-Instruct-q0f16-MLC";
-/** HF root for the default weights (includes `tensor-cache.json` for WebLLM 0.2+). */
-export declare const LOCAL_MLC_HF_MODEL_REPO = "https://huggingface.co/mlc-ai/SmolLM2-135M-Instruct-q0f16-MLC";
-export declare const LOCAL_MLC_MODEL_LIB_WASM_NAME = "SmolLM2-135M-Instruct-q0f16-ctx4k_cs1k-webgpu.wasm";
+export declare 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 declare 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 declare const wrapForArctic: (text: string) => string;
+/**
+ * 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 declare const BE_PARITY: {
+    /** Final payload size cap (BE: `top_k_words`). */
+    readonly TOP_K_WORDS: 2000;
+    /** L2 (domain) words admitted unconditionally before pooling (BE: `reserved_l2_slots`). */
+    readonly RESERVED_L2_SLOTS: 500;
+    /** Log-space additive bias favouring L2 over L3 in the pool (BE: `l2_bias`). */
+    readonly L2_BIAS: 1;
+    /** Drop words below this probability from the final payload (BE: `> 0.00001`). */
+    readonly 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.
+     */
+    readonly 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.
+     */
+    readonly MAX_CONTEXT_WORDS: 100;
+};
+/**
+ * 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 declare const computeBePayload: (rawLogits: Float32Array, prefixMap: Map<number, string[]>, domainWords: Set<string>, 
 /**
- * Original target repo (add-basics fine-tune). **Not compatible with WebLLM 0.2.x** (no `tensor-cache.json`).
- * @see module doc above
+ * 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.
  */
-export declare const HUGGINGFACE_TB_SMOLLM_ADD_BASICS_REPO = "https://huggingface.co/HuggingFaceTB/smollm-135M-instruct-add-basics-q0f16-MLC";
+validTokenIds?: number[]) => Record<string, number>;
 /**
  * Create a local slow-lane client powered by MLC WebLLM.
  *

package/dist/types/pm-plugins/text-predictor.d.ts

@@ -84,5 +84,5 @@ export declare const loadVectorsAsync: (options?: {
     getBinaryUrl?: () => Promise<string>;
 }) => Promise<void>;
 export declare const initVectors: (store: VectorStore) => void;
-export declare const loadDefaultVocabulary: () => void;
+export declare const loadDefaultVocabulary: () => Promise<void>;
 export {};

package/dist/types-ts4.5/pm-plugins/local-slow-lane-client.d.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)
@@ -72,15 +81,93 @@ export interface LocalSlowLaneClient {
     setLmLogits: (logits: Record<string, number> | null) => void;
     updateContext: (text: string) => void;
 }
-export declare const LOCAL_MLC_MODEL_ID = "SmolLM2-135M-Instruct-q0f16-MLC";
-/** HF root for the default weights (includes `tensor-cache.json` for WebLLM 0.2+). */
-export declare const LOCAL_MLC_HF_MODEL_REPO = "https://huggingface.co/mlc-ai/SmolLM2-135M-Instruct-q0f16-MLC";
-export declare const LOCAL_MLC_MODEL_LIB_WASM_NAME = "SmolLM2-135M-Instruct-q0f16-ctx4k_cs1k-webgpu.wasm";
+export declare 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 declare 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 declare const wrapForArctic: (text: string) => string;
+/**
+ * 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 declare const BE_PARITY: {
+    /** Final payload size cap (BE: `top_k_words`). */
+    readonly TOP_K_WORDS: 2000;
+    /** L2 (domain) words admitted unconditionally before pooling (BE: `reserved_l2_slots`). */
+    readonly RESERVED_L2_SLOTS: 500;
+    /** Log-space additive bias favouring L2 over L3 in the pool (BE: `l2_bias`). */
+    readonly L2_BIAS: 1;
+    /** Drop words below this probability from the final payload (BE: `> 0.00001`). */
+    readonly 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.
+     */
+    readonly 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.
+     */
+    readonly MAX_CONTEXT_WORDS: 100;
+};
+/**
+ * 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 declare const computeBePayload: (rawLogits: Float32Array, prefixMap: Map<number, string[]>, domainWords: Set<string>, 
 /**
- * Original target repo (add-basics fine-tune). **Not compatible with WebLLM 0.2.x** (no `tensor-cache.json`).
- * @see module doc above
+ * 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.
  */
-export declare const HUGGINGFACE_TB_SMOLLM_ADD_BASICS_REPO = "https://huggingface.co/HuggingFaceTB/smollm-135M-instruct-add-basics-q0f16-MLC";
+validTokenIds?: number[]) => Record<string, number>;
 /**
  * Create a local slow-lane client powered by MLC WebLLM.
  *

package/dist/types-ts4.5/pm-plugins/text-predictor.d.ts

@@ -84,5 +84,5 @@ export declare const loadVectorsAsync: (options?: {
     getBinaryUrl?: () => Promise<string>;
 }) => Promise<void>;
 export declare const initVectors: (store: VectorStore) => void;
-export declare const loadDefaultVocabulary: () => void;
+export declare const loadDefaultVocabulary: () => Promise<void>;
 export {};

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@atlaskit/editor-plugin-autocomplete",
-	"version": "3.0.0",
+	"version": "3.2.0",
 	"description": "Client-side text autocomplete plugin for @atlaskit/editor-core",
 	"author": "Atlassian Pty Ltd",
 	"license": "Apache-2.0",
@@ -35,7 +35,7 @@
 		"wink-nlp": "^2.4.0"
 	},
 	"peerDependencies": {
-		"@atlaskit/editor-common": "^115.0.0",
+		"@atlaskit/editor-common": "^115.6.0",
 		"@atlaskit/editor-plugin-analytics": "^11.0.0",
 		"react": "^18.2.0"
 	},