change(web): add lexical weighting to prediction search

web/src/engine/predictive-text/worker-thread/src/main/correction/correction-result-mapping.ts

@@ -27,10 +27,16 @@ export interface CorrectionResultMapping<ResultType> {
   readonly matchedResult: Readonly<ResultType>;
 
   /**
-   * Gets the "total cost" of the edge, which should be considered as the
+   * Gets the "correction cost" of the edge, which should be considered as the
    * negative log-likelihood of the input path taken to reach the node
    * multiplied by the 'probability' induced by needed Damerau-Levenshtein edits
    * to the resulting output.
    */
-  readonly totalCost: number;
+  readonly correctionCost: number;
+
+  /**
+   * The "total cost" of the edge - comprised of both the correction cost and the
+   * prediction cost based on the model's frequency data for the word.
+   */
+  readonly currentCost: number;
 }

web/src/engine/predictive-text/worker-thread/src/main/correction/correction-searchable.ts

@@ -29,10 +29,14 @@ type CompleteSearchPath<MappingType> = {
 
 export type PathResult<MappingType> = NullPath | IntermediateSearchPath | CompleteSearchPath<MappingType>;
 
-export function CORRECTION_QUEUE_COMPARATOR<T extends {currentCost: number}>(a: T, b: T) {
+export function PREDICTION_QUEUE_COMPARATOR<T extends {currentCost: number}>(a: T, b: T) {
   return a.currentCost - b.currentCost;
 }
 
+export function CORRECTION_QUEUE_COMPARATOR<T extends {correctionCost: number}>(a: T, b: T) {
+  return a.correctionCost - b.correctionCost;
+}
+
 /**
  * Represents objects that support correction search via the `getBestMatches`
  * method, providing metadata relative to optimizing the search process for

web/src/engine/predictive-text/worker-thread/src/main/correction/distance-modeler.ts

@@ -4,7 +4,7 @@ import { PriorityQueue } from 'keyman/common/web-utils';
 import { LexicalModelTypes } from '@keymanapp/common-types';
 
 import { ClassicalDistanceCalculation } from './classical-calculation.js';
-import { CORRECTION_QUEUE_COMPARATOR, CorrectionSearchable } from './correction-searchable.js';
+import { PREDICTION_QUEUE_COMPARATOR, CorrectionSearchable } from './correction-searchable.js';
 import { CorrectionResultMapping } from './correction-result-mapping.js';
 import { ExecutionTimer, STANDARD_TIME_BETWEEN_DEFERS } from './execution-timer.js';
 import { SearchQuotientNode } from './search-quotient-node.js';
@@ -262,7 +262,7 @@ export class SearchNode {
    * The correction search evaluates Nodes in cost-ascending order based on this property's
    * return value.
    */
-  get currentCost(): number {
+  get correctionCost(): number {
     // - We reintrepret 'known cost' as a psuedo-probability.
     //   - Noting that 1/e = 0.367879441, an edit-distance cost of 1 may be intepreted as -ln(1/e) - a log-space 'likelihood'.
     //     - Not exactly normalized, though.
@@ -278,6 +278,14 @@ export class SearchNode {
     return EDIT_DISTANCE_COST_SCALE * this.editCount + this.inputSamplingCost;
   }
 
+  get predictionCost(): number {
+    return -Math.log(this.currentTraversal.p);
+  }
+
+  get currentCost(): number {
+    return this.correctionCost + this.predictionCost;
+  }
+
   addEdit() {
     this.addedEditCost++;
   }
@@ -625,15 +633,15 @@ export async function *getBestMatches<
   // If no filter function is provided, default to one that always returns true.
   filter ??= () => true;
 
-  let spaceQueue = new PriorityQueue<Correctable>(CORRECTION_QUEUE_COMPARATOR);
+  let spaceQueue = new PriorityQueue<Correctable>(PREDICTION_QUEUE_COMPARATOR);
 
   // Stage 1 - if we already have extracted results, build a queue just for them
   // and iterate over it first.
   //
   // Does not get any results that another iterator pulls up after this is
   // created - and those results won't come up later in stage 2, either.  Only
   // intended for restarting a search, not searching twice in parallel.
-  const priorResultsQueue = new PriorityQueue<ResultMapping>((a, b) => a.totalCost - b.totalCost);
+  const priorResultsQueue = new PriorityQueue<ResultMapping>(PREDICTION_QUEUE_COMPARATOR);
   priorResultsQueue.enqueueAll(searchModules.map((space) => space.previousResults).flat());
 
   // With potential prior results re-queued, NOW enqueue.  (Not before - the heap may reheapify!)
@@ -642,7 +650,7 @@ export async function *getBestMatches<
   // Stage 2:  the fun part; actually searching!
   do {
     const entry: ResultMapping = timer.time(() => {
-      if((priorResultsQueue.peek()?.totalCost ?? Number.POSITIVE_INFINITY) <= spaceQueue.peek().currentCost) {
+      if((priorResultsQueue.peek()?.currentCost ?? Number.POSITIVE_INFINITY) <= spaceQueue.peek().currentCost) {
         const result = priorResultsQueue.dequeue();
 
         // There's no guarantee that the filter closure is the same instance as
@@ -669,7 +677,7 @@ export async function *getBestMatches<
       let lowestCostSource = spaceQueue.dequeue();
       const newResult = lowestCostSource.handleNextNode();
       spaceQueue.enqueue(lowestCostSource);
-      spaceQueue = new PriorityQueue(CORRECTION_QUEUE_COMPARATOR, spaceQueue.toArray());
+      spaceQueue = new PriorityQueue(PREDICTION_QUEUE_COMPARATOR, spaceQueue.toArray());
 
       if(newResult.type == 'none') {
         return null;

web/src/engine/predictive-text/worker-thread/src/main/correction/legacy-quotient-root.ts

@@ -1,7 +1,7 @@
 import { PriorityQueue } from 'keyman/common/web-utils';
 import { LexicalModelTypes } from '@keymanapp/common-types';
 
-import { CORRECTION_QUEUE_COMPARATOR, PathResult } from './correction-searchable.js';
+import { PREDICTION_QUEUE_COMPARATOR, PathResult } from './correction-searchable.js';
 import { SearchQuotientNode } from './search-quotient-node.js';
 import { SearchQuotientRoot } from './search-quotient-root.js';
 import { SearchNode } from './distance-modeler.js';
@@ -10,7 +10,7 @@ import LexicalModel = LexicalModelTypes.LexicalModel;
 import { TokenResultMapping } from './token-result-mapping.js';
 
 export class LegacyQuotientRoot extends SearchQuotientRoot {
-  private selectionQueue: PriorityQueue<SearchNode> = new PriorityQueue(CORRECTION_QUEUE_COMPARATOR);
+  private selectionQueue: PriorityQueue<SearchNode> = new PriorityQueue(PREDICTION_QUEUE_COMPARATOR);
   private processed: SearchNode[] = [];
 
   constructor(model: LexicalModel) {

web/src/engine/predictive-text/worker-thread/src/main/correction/legacy-quotient-spur.ts

@@ -11,7 +11,7 @@
 import { LexicalModelTypes } from '@keymanapp/common-types';
 import { KMWString, PriorityQueue } from 'keyman/common/web-utils';
 
-import { CORRECTION_QUEUE_COMPARATOR, PathResult } from './correction-searchable.js';
+import { PREDICTION_QUEUE_COMPARATOR, PathResult } from './correction-searchable.js';
 import { SearchNode } from './distance-modeler.js';
 import { SearchQuotientNode, PathInputProperties } from './search-quotient-node.js';
 import { SearchQuotientSpur } from './search-quotient-spur.js';
@@ -24,7 +24,7 @@ import Transform = LexicalModelTypes.Transform;
 // The set of search spaces corresponding to the same 'context' for search.
 // Whenever a wordbreak boundary is crossed, a new instance should be made.
 export class LegacyQuotientSpur extends SearchQuotientSpur {
-  private transposeQueue: PriorityQueue<SearchNode> = new PriorityQueue(CORRECTION_QUEUE_COMPARATOR);
+  private transposeQueue: PriorityQueue<SearchNode> = new PriorityQueue(PREDICTION_QUEUE_COMPARATOR);
   private incomingTransposeRootNodes: TokenResultMapping[] = [];
 
   public readonly insertLength: number;

web/src/engine/predictive-text/worker-thread/src/main/correction/search-quotient-cluster.ts

@@ -11,7 +11,7 @@
 import { PriorityQueue } from 'keyman/common/web-utils';
 import { LexicalModelTypes } from '@keymanapp/common-types';
 
-import { CORRECTION_QUEUE_COMPARATOR, PathResult } from './correction-searchable.js';
+import { PREDICTION_QUEUE_COMPARATOR, PathResult } from './correction-searchable.js';
 import { LegacyQuotientRoot } from './legacy-quotient-root.js';
 import { generateSpaceSeed, InputSegment, SearchQuotientNode } from './search-quotient-node.js';
 import { SearchQuotientSpur } from './search-quotient-spur.js';
@@ -20,7 +20,7 @@ import { TokenResultMapping } from './token-result-mapping.js';
 // The set of search spaces corresponding to the same 'context' for search.
 // Whenever a wordbreak boundary is crossed, a new instance should be made.
 export class SearchQuotientCluster extends SearchQuotientNode {
-  private selectionQueue: PriorityQueue<SearchQuotientNode> = new PriorityQueue(CORRECTION_QUEUE_COMPARATOR);
+  private selectionQueue: PriorityQueue<SearchQuotientNode> = new PriorityQueue(PREDICTION_QUEUE_COMPARATOR);
   readonly spaceId: number;
 
   // We use an array and not a PriorityQueue b/c batch-heapifying at a single
@@ -102,7 +102,7 @@ export class SearchQuotientCluster extends SearchQuotientNode {
     entries.forEach((path) => path.increaseMaxEditDistance());
 
     // Since we just modified the stored instances, and the costs may have shifted, we need to re-heapify.
-    this.selectionQueue = new PriorityQueue<SearchQuotientNode>(CORRECTION_QUEUE_COMPARATOR, entries.slice());
+    this.selectionQueue = new PriorityQueue<SearchQuotientNode>(PREDICTION_QUEUE_COMPARATOR, entries.slice());
   }
 
   /**
@@ -131,7 +131,7 @@ export class SearchQuotientCluster extends SearchQuotientNode {
     const bestPath = this.selectionQueue.dequeue();
     const baseResult = bestPath.handleNextNode();
     this.selectionQueue.enqueue(bestPath);
-    this.selectionQueue = new PriorityQueue(CORRECTION_QUEUE_COMPARATOR, this.selectionQueue.toArray());
+    this.selectionQueue = new PriorityQueue(PREDICTION_QUEUE_COMPARATOR, this.selectionQueue.toArray());
 
     let finalResult = baseResult;
     if(baseResult.type == 'complete') {

web/src/engine/predictive-text/worker-thread/src/main/correction/search-quotient-spur.ts

@@ -12,7 +12,7 @@ import { KMWString, PriorityQueue } from 'keyman/common/web-utils';
 import { LexicalModelTypes } from '@keymanapp/common-types';
 import { buildMergedTransform } from '@keymanapp/models-templates';
 
-import { CORRECTION_QUEUE_COMPARATOR, PathResult } from './correction-searchable.js';
+import { PREDICTION_QUEUE_COMPARATOR, PathResult } from './correction-searchable.js';
 import { EDIT_DISTANCE_COST_SCALE, SearchNode } from './distance-modeler.js';
 import { generateSpaceSeed, InputSegment, PathInputProperties, SearchQuotientNode } from './search-quotient-node.js';
 import { generateSubsetId } from './tokenization-subsets.js';
@@ -34,7 +34,7 @@ export const MAX_EDIT_THRESHOLD_FACTOR = 2.5;
 // The set of search spaces corresponding to the same 'context' for search.
 // Whenever a wordbreak boundary is crossed, a new instance should be made.
 export abstract class SearchQuotientSpur extends SearchQuotientNode {
-  private selectionQueue: PriorityQueue<SearchNode> = new PriorityQueue(CORRECTION_QUEUE_COMPARATOR);
+  private selectionQueue: PriorityQueue<SearchNode> = new PriorityQueue(PREDICTION_QUEUE_COMPARATOR);
 
   /**
    * Holds all incoming Nodes generated from a parent `SearchSpace` that have not yet been
@@ -152,7 +152,7 @@ export abstract class SearchQuotientSpur extends SearchQuotientNode {
     entries.forEach(function(edge) { edge.calculation = edge.calculation.increaseMaxDistance(); });
 
     // Since we just modified the stored instances, and the costs may have shifted, we need to re-heapify.
-    this.selectionQueue = new PriorityQueue<SearchNode>(CORRECTION_QUEUE_COMPARATOR, entries);
+    this.selectionQueue = new PriorityQueue<SearchNode>(PREDICTION_QUEUE_COMPARATOR, entries);
   }
 
   /**
@@ -419,7 +419,7 @@ export abstract class SearchQuotientSpur extends SearchQuotientNode {
     // Allows a little 'wiggle room' + 2 "hard" edits.
     // Can be important if needed characters don't actually exist on the keyboard
     // ... or even just not the then-current layer of the keyboard.
-    if(currentNode.currentCost > this.lowestPossibleSingleCost + MAX_EDIT_THRESHOLD_FACTOR * EDIT_DISTANCE_COST_SCALE) {
+    if(currentNode.correctionCost > this.lowestPossibleSingleCost + MAX_EDIT_THRESHOLD_FACTOR * EDIT_DISTANCE_COST_SCALE) {
       return unmatchedResult;
     }
 

web/src/engine/predictive-text/worker-thread/src/main/correction/token-result-mapping.ts

@@ -33,8 +33,8 @@ export function initTokenResultFilterer() {
       return false;
     }
 
-    if((priorReturnCosts.get(searchResult.matchString) ?? Number.MAX_VALUE) > searchResult.totalCost) {
-      priorReturnCosts.set(searchResult.matchString, searchResult.totalCost);
+    if((priorReturnCosts.get(searchResult.matchString) ?? Number.MAX_VALUE) > searchResult.correctionCost) {
+      priorReturnCosts.set(searchResult.matchString, searchResult.correctionCost);
 
       return true;
     } else {
@@ -118,6 +118,14 @@ export class TokenResultMapping implements CorrectionResultMapping<SearchNode> {
    * multiplied by the 'probability' induced by needed Damerau-Levenshtein edits
    * to the resulting output.
    */
+  get correctionCost(): number {
+    return this.node.correctionCost;
+  }
+
+  get currentCost(): number {
+    return this.node.currentCost;
+  }
+
   get totalCost(): number {
     return this.node.currentCost;
   }

web/src/engine/predictive-text/worker-thread/src/main/predict-helpers.ts

@@ -532,7 +532,7 @@ export function buildAndMapPredictions(
   transition: ContextTransition,
   tokenization: ContextTokenization,
   // Originally, Readonly<TokenResultMapping> - but we only need these three components here.
-  match: Readonly<{matchString: string, totalCost: number, editCount: number}>,
+  match: Readonly<{matchString: string, correctionCost: number, editCount: number}>,
   costFactor: number
 ): CorrectionPredictionTuple[] {
   const model = transition.final.model;
@@ -545,7 +545,7 @@ export function buildAndMapPredictions(
 
   // --- to move into predictFromCorrections ---
   let correction = match.matchString;
-  let rootCost = match.totalCost;
+  let rootCost = match.correctionCost;
 
   // Replace the existing context with the correction.
   const correctionTransform: Transform = {
@@ -656,7 +656,7 @@ export async function correctAndEnumerate(
 
   // Only run the correction search when corrections are enabled.
   let rawPredictions: CorrectionPredictionTuple[] = [];
-  let bestCorrectionCost: number;
+  let bestTotalCost: number;
   const correctionPredictionMap: Record<string, Distribution<Suggestion>> = {};
   for await(const match of getBestTokenMatches(searchModules, timer)) {
     // Corrections obtained:  now to predict from them!
@@ -700,8 +700,8 @@ export async function correctAndEnumerate(
     const predictions = buildAndMapPredictions(transition, tokenization, match, costFactor);
 
     // Only set 'best correction' cost when a correction ACTUALLY YIELDS predictions.
-    if(predictions.length > 0 && bestCorrectionCost === undefined) {
-      bestCorrectionCost = match.totalCost * costFactor;
+    if(predictions.length > 0 && bestTotalCost === undefined) {
+      bestTotalCost = match.totalCost * costFactor;
     }
 
     // If we're getting the same prediction again, it's lower-cost.  Update!
@@ -714,7 +714,7 @@ export async function correctAndEnumerate(
 
     rawPredictions = rawPredictions.concat(predictions);
 
-    if(shouldStopSearchingEarly(bestCorrectionCost, match.totalCost, rawPredictions)) {
+    if(shouldStopSearchingEarly(bestTotalCost, match.totalCost, rawPredictions)) {
       break;
     }
   }
@@ -738,20 +738,15 @@ export function shouldStopSearchingEarly(
     return true;
     // If enough have been found, we're safe to terminate earlier.
   } else if(rawPredictions.length >= ModelCompositor.MAX_SUGGESTIONS) {
-    if(currentCorrectionCost >= bestCorrectionCost + CORRECTION_SEARCH_THRESHOLDS.REPLACEMENT_SEARCH_THRESHOLD) {
-      // Very useful for stopping 'sooner' when words reach a sufficient length.
+    // Sort the prediction list; we need them in descending probability order
+    // for the next check.
+    rawPredictions.sort((a, b) => b.totalProb - a.totalProb);
+
+    // If the best result at the current state of the search fails to beat the worst
+    // pending suggestion from previous tiers, assume all further corrections will
+    // similarly fail to win; terminate the search-loop.
+    if(rawPredictions[ModelCompositor.MAX_SUGGESTIONS-1].totalProb > Math.exp(-currentCorrectionCost)) {
       return true;
-    } else {
-      // Sort the prediction list; we need them in descending probability order
-      // for the next check.
-      rawPredictions.sort((a, b) => b.totalProb - a.totalProb);
-
-      // If the best result at the current state of the search fails to beat the worst
-      // pending suggestion from previous tiers, assume all further corrections will
-      // similarly fail to win; terminate the search-loop.
-      if(rawPredictions[ModelCompositor.MAX_SUGGESTIONS-1].totalProb > Math.exp(-currentCorrectionCost)) {
-        return true;
-      }
     }
   }