@atlaskit/editor-plugin-autocomplete 0.1.0 → 0.3.0

package/dist/esm/pm-plugins/autocomplete-plugin.js

@@ -0,0 +1,362 @@
+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; }
+// url: prefix is an Atlaspack/Parcel directive that resolves this file as an emitted
+// asset URL (content-hashed). For rspack, this is handled via staticAssetsLoader.
+// eslint-disable-next-line @repo/internal/import/no-unresolved
+import wordVectorsUrl from 'url:./data/word-vectors_10k.bin';
+import { SafePlugin } from '@atlaskit/editor-common/safe-plugin';
+import { keydownHandler } from '@atlaskit/editor-prosemirror/keymap';
+import { PluginKey } from '@atlaskit/editor-prosemirror/state';
+import { DecorationSet } from '@atlaskit/editor-prosemirror/view';
+import { createGhostTextDecorationSet } from './ghost-text-decoration';
+import { createSlowLaneClient, setDefaultSlowLaneClient, isWordBoundary } from './slow-lane-client';
+import { predict, loadDefaultVocabulary, loadVectorsAsync, incrementSessionFreq, ingestDocumentPage } from './text-predictor';
+var SLOW_LANE_ENDPOINT = '/gateway/api/v1/autocomplete/typeahead-encodings';
+export var autocompletePluginKey = new PluginKey('autocomplete');
+var DEBOUNCE_MS = 150;
+var createInitialState = function createInitialState() {
+  return {
+    ghostText: '',
+    ghostPosition: -1,
+    decorationSet: DecorationSet.empty
+  };
+};
+
+/**
+ * Extract text content before the cursor from the current document.
+ * Returns the last ~200 characters for context.
+ */
+var getTextBeforeCursor = function getTextBeforeCursor(state) {
+  var $from = state.selection.$from;
+  var maxChars = 200;
+
+  // 1. Get the perfectly flattened text of the current block up to the cursor
+  var blockNode = $from.parent;
+  var offsetInBlock = $from.parentOffset;
+  var blockText = blockNode.textContent.slice(0, offsetInBlock);
+  if (blockText.length >= maxChars) {
+    return blockText.slice(-maxChars);
+  }
+  var fullText = blockText;
+
+  // 2. Walk backwards through previous blocks
+  var depth = $from.depth - 1;
+  while (fullText.length < maxChars && depth >= 0) {
+    var parentNode = $from.node(depth);
+    var indexInParent = $from.index(depth);
+    for (var i = indexInParent - 1; i >= 0 && fullText.length < maxChars; i--) {
+      var sibling = parentNode.child(i);
+      var siblingText = sibling.textContent;
+      fullText = siblingText + '\n' + fullText;
+    }
+    depth--;
+  }
+  return fullText.slice(-maxChars);
+};
+
+/**
+ * Set the autocomplete state via a transaction metadata.
+ */
+var setAutocompleteMeta = function setAutocompleteMeta(tr, meta) {
+  return tr.setMeta(autocompletePluginKey, meta);
+};
+
+/**
+ * Apply a ghost text suggestion to the editor state.
+ */
+var showGhostText = function showGhostText(view, text, position) {
+  var state = view.state,
+    dispatch = view.dispatch;
+  var decorationSet = createGhostTextDecorationSet(state, position, text);
+  var tr = setAutocompleteMeta(state.tr, {
+    ghostText: text,
+    ghostPosition: position,
+    decorationSet: decorationSet
+  });
+  dispatch(tr);
+};
+
+/**
+ * Clear the current ghost text from the editor.
+ */
+var clearGhostText = function clearGhostText(state, dispatch) {
+  var pluginState = autocompletePluginKey.getState(state);
+  if (!pluginState || !pluginState.ghostText) {
+    return false;
+  }
+  if (dispatch) {
+    var tr = setAutocompleteMeta(state.tr, {
+      ghostText: '',
+      ghostPosition: -1,
+      decorationSet: DecorationSet.empty
+    });
+    dispatch(tr);
+  }
+  return true;
+};
+
+/**
+ * Accept the current ghost text suggestion and insert it into the document.
+ */
+var acceptGhostText = function acceptGhostText(state, dispatch) {
+  var pluginState = autocompletePluginKey.getState(state);
+  if (!pluginState || !pluginState.ghostText) {
+    return false;
+  }
+  if (dispatch) {
+    var ghostText = pluginState.ghostText,
+      ghostPosition = pluginState.ghostPosition;
+    var tr = state.tr.insertText(ghostText, ghostPosition);
+    tr = setAutocompleteMeta(tr, {
+      ghostText: '',
+      ghostPosition: -1,
+      decorationSet: DecorationSet.empty
+    });
+    dispatch(tr);
+  }
+  return true;
+};
+
+/**
+ * Context provided to the autocomplete plugin on first editor focus.
+ * Text fields are selectively ingested to boost word-frequency scoring for
+ * predictions, giving words already present in the document/thread an L1
+ * priority boost.
+ */
+
+/**
+ * Build the text payload for the slow-lane request by prepending any available
+ * comment context ahead of the live document text. This gives the backend
+ * model richer context about the thread the user is writing in.
+ */
+var buildSlowLaneText = function buildSlowLaneText(docText, context) {
+  var _context$siblingComme, _context$siblingComme2, _context$siblingComme3;
+  var lines = [];
+  if (context !== null && context !== void 0 && context.parentCommentContent) {
+    lines.push("comment: ".concat(context.parentCommentContent));
+  }
+  context === null || context === void 0 || (_context$siblingComme = context.siblingCommentsContents) === null || _context$siblingComme === void 0 || _context$siblingComme.forEach(function (sibling, index) {
+    lines.push("reply ".concat(index + 1, ": ").concat(sibling));
+  });
+  var nextReplyNumber = ((_context$siblingComme2 = context === null || context === void 0 || (_context$siblingComme3 = context.siblingCommentsContents) === null || _context$siblingComme3 === void 0 ? void 0 : _context$siblingComme3.length) !== null && _context$siblingComme2 !== void 0 ? _context$siblingComme2 : 0) + 1;
+  lines.push("reply ".concat(nextReplyNumber, ": ").concat(docText));
+  return lines.join('\n');
+};
+export var createAutocompletePlugin = function createAutocompletePlugin(options) {
+  var debounceTimer = null;
+  var hasIngestedPage = false;
+  var resolvedContext;
+  /**
+   * Set after accepting a suggestion so the next doc-change update
+   * skips scheduling a new prediction for the just-inserted text.
+   * Scoped to the factory so multiple editor instances don't share state.
+   */
+  var justAccepted = false;
+
+  /**
+   * Stores the text-before-cursor snapshot at the moment the user dismissed
+   * a suggestion via Escape. While the context remains identical, we suppress
+   * re-showing the same suggestion. Resets to null as soon as the text changes.
+   */
+  var dismissedContext = null;
+  var slowLaneClient = createSlowLaneClient({
+    baseUrl: '',
+    endpoint: SLOW_LANE_ENDPOINT
+  });
+  setDefaultSlowLaneClient(slowLaneClient);
+
+  /**
+   * Schedule a prediction after a short debounce.
+   * Tier 1 predictions are synchronous (<0.1ms) but we still debounce
+   * to avoid unnecessary work on rapid keystrokes.
+   */
+  var schedulePrediction = function schedulePrediction(view) {
+    if (debounceTimer) {
+      clearTimeout(debounceTimer);
+    }
+    debounceTimer = setTimeout(function () {
+      var state = view.state;
+      var selection = state.selection;
+
+      // Only predict for cursor selections (not range selections)
+      if (!selection.empty) {
+        return;
+      }
+      var textBefore = getTextBeforeCursor(state);
+
+      // Suppress re-showing the same suggestion the user just dismissed.
+      // Once the text context changes (user types or deletes), this clears automatically.
+      if (textBefore === dismissedContext) {
+        return;
+      }
+      dismissedContext = null;
+
+      // Don't predict if there's not enough context
+      if (textBefore.trim().length < 3) {
+        return;
+      }
+
+      // Tier 1 prediction is synchronous -- no async needed
+      var prediction = predict(textBefore);
+      if (prediction && prediction.length > 0) {
+        showGhostText(view, prediction, selection.from);
+      }
+    }, DEBOUNCE_MS);
+  };
+  var maybeUpdateSessionFrequency = function maybeUpdateSessionFrequency(view, prevState) {
+    var newText = getTextBeforeCursor(view.state);
+    var prevText = getTextBeforeCursor(prevState);
+    if (newText.length <= prevText.length) {
+      return;
+    }
+    var lastChar = newText[newText.length - 1];
+    // eslint-disable-next-line require-unicode-regexp
+    if (!/[\s.,;:!?]/.test(lastChar)) {
+      return;
+    }
+
+    // Only fire if the previous state did not already end on a boundary,
+    // so we don't double-count when multiple boundary chars are inserted.
+    var prevLastChar = prevText[prevText.length - 1];
+    // eslint-disable-next-line require-unicode-regexp
+    if (prevLastChar && /[\s.,;:!?]/.test(prevLastChar)) {
+      return;
+    }
+    var beforeBoundary = newText.slice(0, -1).trimEnd();
+    var lastSpaceIdx = beforeBoundary.lastIndexOf(' ');
+    var completedWord = beforeBoundary.slice(lastSpaceIdx + 1).toLowerCase();
+    if (completedWord.length >= 2) {
+      incrementSessionFreq(completedWord);
+    }
+  };
+  return new SafePlugin({
+    key: autocompletePluginKey,
+    state: {
+      init: function init() {
+        return createInitialState();
+      },
+      apply: function apply(tr, pluginState) {
+        var meta = tr.getMeta(autocompletePluginKey);
+        if (meta) {
+          return _objectSpread(_objectSpread({}, pluginState), meta);
+        }
+
+        // If the document changed, clear the ghost text
+        // (new prediction will be scheduled from view.update)
+        if (tr.docChanged) {
+          return _objectSpread(_objectSpread({}, pluginState), {}, {
+            ghostText: '',
+            ghostPosition: -1,
+            decorationSet: DecorationSet.empty
+          });
+        }
+
+        // If selection changed without doc change, clear ghost text
+        if (tr.selectionSet && pluginState.ghostText) {
+          return _objectSpread(_objectSpread({}, pluginState), {}, {
+            ghostText: '',
+            ghostPosition: -1,
+            decorationSet: DecorationSet.empty
+          });
+        }
+        return pluginState;
+      }
+    },
+    props: {
+      decorations: function decorations(state) {
+        var _pluginState$decorati;
+        var pluginState = autocompletePluginKey.getState(state);
+        return (_pluginState$decorati = pluginState === null || pluginState === void 0 ? void 0 : pluginState.decorationSet) !== null && _pluginState$decorati !== void 0 ? _pluginState$decorati : DecorationSet.empty;
+      },
+      handleKeyDown: keydownHandler({
+        Tab: function Tab(state, dispatch) {
+          var accepted = acceptGhostText(state, dispatch);
+          if (accepted) justAccepted = true;
+          return accepted;
+        },
+        ArrowRight: function ArrowRight(state, dispatch) {
+          var accepted = acceptGhostText(state, dispatch);
+          if (accepted) justAccepted = true;
+          return accepted;
+        },
+        Escape: function Escape(state, dispatch) {
+          var didClear = clearGhostText(state, dispatch);
+          if (didClear) {
+            dismissedContext = getTextBeforeCursor(state);
+          }
+          return didClear;
+        }
+      }),
+      handleDOMEvents: {
+        blur: function blur(view) {
+          var pluginState = autocompletePluginKey.getState(view.state);
+          if (pluginState !== null && pluginState !== void 0 && pluginState.ghostText) {
+            clearGhostText(view.state, view.dispatch);
+          }
+          return false;
+        },
+        focus: function focus() {
+          loadDefaultVocabulary();
+          loadVectorsAsync({
+            vectorsUrl: wordVectorsUrl
+          }).catch(function () {});
+          if (!hasIngestedPage) {
+            hasIngestedPage = true;
+            if (options !== null && options !== void 0 && options.getContext) {
+              options.getContext().then(function (context) {
+                var _context$siblingComme4;
+                if (!context) {
+                  return;
+                }
+                resolvedContext = context;
+                if (context.fullPageContent) {
+                  ingestDocumentPage(context.fullPageContent);
+                }
+                if (context.parentCommentContent) {
+                  ingestDocumentPage(context.parentCommentContent);
+                }
+                (_context$siblingComme4 = context.siblingCommentsContents) === null || _context$siblingComme4 === void 0 || _context$siblingComme4.forEach(ingestDocumentPage);
+              }).catch(function () {});
+            }
+          }
+          return false;
+        }
+      }
+    },
+    view: function view() {
+      return {
+        update: function update(view, prevState) {
+          if (!prevState.doc.eq(view.state.doc)) {
+            if (justAccepted) {
+              justAccepted = false;
+
+              // ✨ THE FIX: Memorize the text state right after acceptance.
+              // Any follow-up transactions will hit the 'dismissedContext'
+              // block and abort until the user actually types a new character!
+              dismissedContext = getTextBeforeCursor(view.state);
+
+              // Also clear any pending debounce timers from before the acceptance
+              if (debounceTimer) {
+                clearTimeout(debounceTimer);
+              }
+              return;
+            }
+            maybeUpdateSessionFrequency(view, prevState);
+            var textBefore = getTextBeforeCursor(view.state);
+            if (isWordBoundary(textBefore)) {
+              slowLaneClient.updateContext(buildSlowLaneText(view.state.doc.textContent, resolvedContext));
+            }
+            schedulePrediction(view);
+          }
+        },
+        destroy: function destroy() {
+          if (debounceTimer) {
+            clearTimeout(debounceTimer);
+          }
+          setDefaultSlowLaneClient(null);
+        }
+      };
+    }
+  });
+};

package/dist/esm/pm-plugins/ghost-text-decoration.js

@@ -0,0 +1,33 @@
+import { Decoration, DecorationSet } from '@atlaskit/editor-prosemirror/view';
+var GHOST_TEXT_CLASS = 'autocomplete-ghost-text';
+
+/**
+ * Creates a DecorationSet containing a ghost text widget at the given position.
+ * The ghost text is rendered as a styled <span> that appears after the cursor.
+ */
+export var createGhostTextDecorationSet = function createGhostTextDecorationSet(state, position, text) {
+  if (!text) {
+    return DecorationSet.empty;
+  }
+  var decoration = Decoration.widget(position, function () {
+    var container = document.createElement('span');
+    container.className = GHOST_TEXT_CLASS;
+    container.setAttribute('data-autocomplete-ghost', 'true');
+    container.style.color = '#999';
+    container.style.opacity = '0.6';
+    container.style.pointerEvents = 'none';
+    container.style.userSelect = 'none';
+    container.style.fontStyle = 'italic';
+    // U+200B (Zero Width Space) gives the browser a line-break opportunity
+    // immediately before the ghost text. This ensures the typed text before
+    // the span is never pushed to the next line by the ghost text's width —
+    // only the ghost text itself will wrap if it doesn't fit.
+    container.textContent = "\u200B" + text;
+    return container;
+  }, {
+    side: 1,
+    // Render after content at this position
+    key: 'autocomplete-ghost-text'
+  });
+  return DecorationSet.create(state.doc, [decoration]);
+};

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

@@ -0,0 +1,252 @@
+import _slicedToArray from "@babel/runtime/helpers/slicedToArray";
+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; }
+/**
+ * Scoring Pipeline: Stage 1 (Semantic + Frequency), Grammar Filter, Stage 2 (LM Re-ranking).
+ *
+ * Operates synchronously on pre-loaded data. Each stage gracefully degrades
+ * when its required data isn't available (cold → warm → full warm).
+ */
+
+import posTagsData from './data/combined_l2_l3_pos_tags.json';
+import ghostPosTagsData from './data/ghost_pos_tags.json';
+import grammarTransitionsData from './data/grammar_transitions_10k.json';
+
+// ─── Types ──────────────────────────────────────────────────
+
+/** Metadata returned by the grammar filter for debug logging in the caller. */
+
+// ─── Scoring Constants ──────────────────────────────────────
+
+var ALPHA = 0.5;
+var BETA = 0.5;
+var NEUTRAL_SCORE = 0.5;
+var STAGE1_WEIGHT = 0.6;
+var STAGE2_WEIGHT = 0.4;
+var MIN_STAGE1_SCORE = 0.35;
+var L1_SESSION_CAP = 1.2;
+
+// ─── Grammar Data (loaded once on import) ───────────────────
+
+var posTags = new Map(Object.entries(posTagsData));
+var grammarTransitions = grammarTransitionsData;
+
+/**
+ * Precomputed map from each POS tag to the set of allowed next POS tags.
+ * Built once at module load from grammarTransitions so applyGrammarFilter
+ * never re-iterates the transition rules per call.
+ */
+var precomputedAllowedByPos = new Map(Object.entries(grammarTransitions.transitions).map(function (_ref) {
+  var _ref2 = _slicedToArray(_ref, 2),
+    pos = _ref2[0],
+    rule = _ref2[1];
+  return [pos, new Set(rule.allowed)];
+}));
+
+// ─── Math ───────────────────────────────────────────────────
+
+function cosineSimilarity(a, b) {
+  var dot = 0;
+  var normA = 0;
+  var normB = 0;
+  for (var i = 0; i < a.length; i++) {
+    dot += a[i] * b[i];
+    normA += a[i] * a[i];
+    normB += b[i] * b[i];
+  }
+  var dNormA = Math.sqrt(normA);
+  var dNormB = Math.sqrt(normB);
+  if (dNormA === 0 || dNormB === 0) {
+    return NEUTRAL_SCORE;
+  }
+  return (1 + dot / (dNormA * dNormB)) / 2;
+}
+
+// ─── Stage 1: Semantic + Frequency ──────────────────────────
+
+function scoreStage1(candidate, contextVector, getWordVector, maxTenantFreq) {
+  // 1. Calculate Base Global Score (Normalized Log)
+  var maxPossibleLog = Math.log10(maxTenantFreq + 1);
+
+  // Diversity Adjustment
+  var diversityRaw = (Math.log10(candidate.tenantFreq + 1) * 0.50 + Math.log10(candidate.docFreq + 1) * 0.25 + Math.log10(candidate.authorFreq + 1) * 0.25) / maxPossibleLog;
+  var sessionMultiplier = candidate.sessionFreq > 0 ? 1 + Math.log10(candidate.sessionFreq + 1) * 2.5 : 1;
+
+  // Apply multiplier; capped at L1_SESSION_CAP (default 1.2) to prevent excessive over-indexing
+  var freqScore = Math.min(diversityRaw * sessionMultiplier, L1_SESSION_CAP);
+
+  // 3. Semantic Scoring
+  var semanticScore = NEUTRAL_SCORE;
+  if (contextVector) {
+    var wordVec = getWordVector(candidate.word);
+    semanticScore = wordVec ? cosineSimilarity(contextVector, wordVec) : NEUTRAL_SCORE;
+  }
+  return {
+    semanticScore: semanticScore,
+    freqScore: freqScore,
+    stage1Score: ALPHA * semanticScore + BETA * freqScore
+  };
+}
+
+// ─── Grammar Filter ─────────────────────────────────────────
+
+/**
+ * GHOST POS DICTIONARY
+ * A hardcoded mapping of common structural English words that were stripped 
+ * from the main domain vocabulary. This allows the grammar filter to understand 
+ * context without suggesting these words to the user.
+ */
+var ghostPosTags = ghostPosTagsData;
+function applyGrammarFilter(candidates, previousWord) {
+  if (!previousWord) return {
+    filtered: candidates,
+    grammarMeta: null
+  };
+  var lowerPrev = previousWord.toLowerCase();
+  var prevTags = ghostPosTags[lowerPrev] || posTags.get(lowerPrev);
+  if (!prevTags || prevTags.length === 0) {
+    return {
+      filtered: candidates,
+      grammarMeta: null
+    };
+  }
+  var allowedNextTags;
+  if (prevTags.length === 1) {
+    var _precomputedAllowedBy;
+    // Common case: single POS tag — reuse the precomputed Set directly (no allocation)
+    allowedNextTags = (_precomputedAllowedBy = precomputedAllowedByPos.get(prevTags[0])) !== null && _precomputedAllowedBy !== void 0 ? _precomputedAllowedBy : new Set();
+  } else {
+    allowedNextTags = new Set();
+    var _iterator = _createForOfIteratorHelper(prevTags),
+      _step;
+    try {
+      for (_iterator.s(); !(_step = _iterator.n()).done;) {
+        var pt = _step.value;
+        var allowed = precomputedAllowedByPos.get(pt);
+        if (allowed) allowed.forEach(function (tag) {
+          return allowedNextTags.add(tag);
+        });
+      }
+    } catch (err) {
+      _iterator.e(err);
+    } finally {
+      _iterator.f();
+    }
+  }
+  var filtered = [];
+  var dropped = [];
+  var _iterator2 = _createForOfIteratorHelper(candidates),
+    _step2;
+  try {
+    for (_iterator2.s(); !(_step2 = _iterator2.n()).done;) {
+      var entry = _step2.value;
+      var candidateTags = posTags.get(entry.candidate.word.toLowerCase());
+
+      // If candidate has no tags (unknown word), let it pass to be safe
+      if (!candidateTags || candidateTags.length === 0) {
+        filtered.push(entry);
+        continue;
+      }
+      if (candidateTags.some(function (ct) {
+        return allowedNextTags.has(ct);
+      })) {
+        filtered.push(entry);
+      } else {
+        dropped.push(entry.candidate.word);
+      }
+    }
+  } catch (err) {
+    _iterator2.e(err);
+  } finally {
+    _iterator2.f();
+  }
+  var finalFiltered = filtered.length > 0 ? filtered : candidates;
+  return {
+    filtered: finalFiltered,
+    grammarMeta: {
+      prevWord: lowerPrev,
+      prevTags: prevTags,
+      before: candidates.length,
+      after: finalFiltered.length,
+      dropped: filtered.length > 0 ? dropped : []
+    }
+  };
+}
+
+// ─── Stage 2: LM Re-ranking ────────────────────────────────
+function getLmScore(word, lmLogits) {
+  if (!lmLogits) return 0;
+
+  // Look up the word directly! No more tokens.
+  var val = lmLogits[word.toLowerCase()];
+  if (typeof val === 'number') {
+    return val;
+  }
+  return 0;
+}
+
+// ─── Public API ─────────────────────────────────────────────
+
+export function rankCandidates(candidates, contextVector, getWordVector, lmLogits, maxTenantFreq, previousWord) {
+  // Stage 1
+  var stage1Results = candidates.map(function (candidate) {
+    var _scoreStage = scoreStage1(candidate, contextVector, getWordVector, maxTenantFreq),
+      semanticScore = _scoreStage.semanticScore,
+      freqScore = _scoreStage.freqScore,
+      stage1Score = _scoreStage.stage1Score;
+    return {
+      candidate: candidate,
+      semanticScore: semanticScore,
+      freqScore: freqScore,
+      stage1Score: stage1Score
+    };
+  });
+  var stage1Survivors = stage1Results.filter(function (entry) {
+    return entry.stage1Score >= MIN_STAGE1_SCORE;
+  });
+
+  // Grammar Filter
+  var _applyGrammarFilter = applyGrammarFilter(stage1Survivors, previousWord),
+    filtered = _applyGrammarFilter.filtered,
+    grammarMeta = _applyGrammarFilter.grammarMeta;
+
+  // Stage 2 + final assembly
+  var lmMax = 0;
+  if (lmLogits && Object.keys(lmLogits).length > 0) {
+    var values = Object.values(lmLogits);
+    lmMax = Math.max.apply(Math, values);
+  }
+  var scored = filtered.map(function (entry) {
+    var lmScore = 0;
+    var finalScore = entry.stage1Score;
+    if (lmLogits && Object.keys(lmLogits).length > 0) {
+      var rawLm = getLmScore(entry.candidate.word, lmLogits);
+      if (rawLm !== 0) {
+        // The word was in the top_k! Score it normally.
+        var logitDiff = Math.log(rawLm) - Math.log(lmMax);
+        lmScore = Math.exp(logitDiff);
+      } else {
+        lmScore = 0.05;
+      }
+      finalScore = STAGE1_WEIGHT * entry.stage1Score + STAGE2_WEIGHT * lmScore;
+    }
+    return {
+      word: entry.candidate.word,
+      freqScore: entry.freqScore,
+      semanticScore: entry.semanticScore,
+      lmScore: lmScore,
+      finalScore: finalScore
+    };
+  });
+  scored.sort(function (a, b) {
+    if (b.finalScore !== a.finalScore) {
+      return b.finalScore - a.finalScore;
+    }
+    return a.word.length - b.word.length;
+  });
+  return {
+    candidates: scored,
+    grammarMeta: grammarMeta
+  };
+}