@atlaskit/editor-plugin-autocomplete 3.1.0 → 3.3.0

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-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/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@atlaskit/editor-plugin-autocomplete",
-	"version": "3.1.0",
+	"version": "3.3.0",
 	"description": "Client-side text autocomplete plugin for @atlaskit/editor-core",
 	"author": "Atlassian Pty Ltd",
 	"license": "Apache-2.0",

package/scripts/gen_first_token_to_words.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""
+Offline build script: generate `first_token_to_words.json`.
+
+This is a ONE-TIME / build-time tool. It is never imported by the plugin and
+never runs in CI. It reproduces the prefix-expansion map that the BE sidecar
+builds in-memory at runtime (`cc-smarts/python-sidecar/src/causal_lm_encoder.py`
+`CausalLMEncoder._ensure_loaded`), so the local (client-only) slow-lane client
+can ship it as a static artifact instead of running a tokenizer in the browser.
+
+The output maps each SmolLM2 first-token id to every L2/L3 vocabulary word whose
+space-prefixed encoding starts with that token. The local client loads it and,
+per inference, spreads the next-token logit mass over whole words (masked softmax
++ prefix expansion) to match the BE's whole-word `lm_logits` payload.
+
+Three details MUST match the BE exactly, or the map is silently wrong:
+  1. Tokenizer = HuggingFaceTB/SmolLM2-135M (base; vocab identical to Instruct).
+  2. Leading space: encode(" " + word) — BPE tokenizes " word" != "word".
+  3. add_special_tokens=False — no BOS/EOS, we want the word's own first token.
+
+Usage:
+    pip install transformers      # torch NOT required (SmolLM2 uses a fast tokenizer)
+    python scripts/gen_first_token_to_words.py
+
+Run it from anywhere — paths are resolved relative to this file's location.
+"""
+
+import os
+import json
+import collections
+
+from transformers import AutoTokenizer
+
+# Ground truth: must match the BE tokenizer (causal_lm_encoder.py line 30).
+TOKENIZER_NAME = "HuggingFaceTB/SmolLM2-135M"
+
+# Resolve data paths relative to this script, so cwd does not matter.
+_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+_DATA_DIR = os.path.join(_SCRIPT_DIR, "..", "src", "pm-plugins", "data")
+L2_PATH = os.path.join(_DATA_DIR, "vocabulary_10k.json")
+L3_PATH = os.path.join(_DATA_DIR, "l3_vocabulary.json")
+OUT_PATH = os.path.join(_DATA_DIR, "first_token_to_words.json")
+
+# Probe words for the post-build sanity check.
+# Must be words that actually appear in vocabulary_10k.json (L2) or l3_vocabulary.json (L3).
+# Common stop words like "the" / "a" are NOT in either vocabulary by design.
+# L2 confirmed: "atlassian", "service", "product", "customer" (vocabulary_10k.json lines 3-18)
+# L3 confirmed: "about", "search", "information", "business"  (l3_vocabulary.json lines 2-15)
+_PROBE_WORDS = ["atlassian", "service", "product", "about", "search"]
+
+
+def load_l2_words(path):
+    """
+    Load the L2 (Atlassian-domain) vocabulary as a set of words.
+
+    The file shape is {"words": {"<word>": {freq, ...}}}, matching the BE's
+    `vocab_data.get("words", {})`. Only the keys are needed.
+
+    :params:
+        path: Absolute path to vocabulary_10k.json
+    :returns:
+        A set of L2 word strings
+    """
+    with open(path, "r") as f:
+        data = json.load(f)
+    return set(data["words"].keys())
+
+
+def load_l3_words(path):
+    """
+    Load the L3 (general English) vocabulary as a list of words.
+
+    The file shape is a flat JSON array of strings, matching the BE's L3 list.
+
+    :params:
+        path: Absolute path to l3_vocabulary.json
+    :returns:
+        A list of L3 word strings
+    """
+    with open(path, "r") as f:
+        return json.load(f)
+
+
+def build_first_token_map(tokenizer, words):
+    """
+    Build the first-token-id -> [words] prefix-expansion map.
+
+    Mirrors the BE loop exactly: each word is encoded with a leading space and
+    no special tokens, and the word is bucketed under its first token id. A set
+    of words is expected so each word is processed once (L2/L3 overlap removed).
+
+    :params:
+        tokenizer: A HuggingFace tokenizer for SmolLM2
+        words: An iterable of unique words (L2 union L3)
+    :returns:
+        A dict mapping int first-token-id to a list of word strings
+    """
+    table = collections.defaultdict(list)
+    for word in words:
+        ids = tokenizer.encode(" " + word, add_special_tokens=False)
+        if ids:
+            table[ids[0]].append(word)
+    return table
+
+
+def verify_map(tokenizer, table, probe_words):
+    """
+    Sanity-check the generated map by confirming probe words land in the right
+    first-token bucket.
+
+    :params:
+        tokenizer: The same SmolLM2 tokenizer used to build the map
+        table: The dict mapping int first-token-id to a list of words
+        probe_words: A list of words expected to be present in the map
+    :returns:
+        None. Raises AssertionError if any probe word is missing or misplaced.
+    """
+    for word in probe_words:
+        ids = tokenizer.encode(" " + word, add_special_tokens=False)
+        assert ids, f"Probe word '{word}' produced no tokens"
+        first_token_id = ids[0]
+        bucket = table.get(first_token_id, [])
+        assert word in bucket, (
+            f"Probe word '{word}' missing under token {first_token_id} "
+            f"(bucket head: {bucket[:5]})"
+        )
+        print(f"  OK: '{word}' -> token {first_token_id} -> {bucket[:5]}...")
+
+
+def main():
+    """
+    Generate first_token_to_words.json from the L2 and L3 vocabularies.
+
+    :params:
+        None
+    :returns:
+        None. Writes the JSON artifact to OUT_PATH and prints a summary.
+    """
+    print(f"[gen] Loading tokenizer: {TOKENIZER_NAME} ...")
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
+
+    print(f"[gen] Loading vocabularies ...")
+    l2_words = load_l2_words(L2_PATH)
+    l3_words = load_l3_words(L3_PATH)
+    all_words = l2_words.union(l3_words)
+
+    print(f"[gen] Building prefix-expansion map for {len(all_words)} words ...")
+    table = build_first_token_map(tokenizer, all_words)
+
+    # JSON object keys must be strings; the FE parses them back with Number(key).
+    out = {str(token_id): words for token_id, words in table.items()}
+    with open(OUT_PATH, "w") as f:
+        json.dump(out, f)
+
+    total_words = sum(len(v) for v in table.values())
+    print(
+        f"[gen] Mapped {len(all_words)} words ({len(l2_words)} L2) "
+        f"-> {len(table)} unique first-tokens ({total_words} word entries)"
+    )
+
+    print(f"[gen] Verifying probe words ...")
+    verify_map(tokenizer, table, _PROBE_WORDS)
+
+    out_abs = os.path.abspath(OUT_PATH)
+    size_kb = os.path.getsize(OUT_PATH) / 1024
+    print(f"[gen] Wrote {out_abs} ({size_kb:.0f} KB)")
+
+
+if __name__ == "__main__":
+    main()