change(web): predict from per-token correction sequence

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

@@ -12,7 +12,7 @@ import { ContextState, determineContextSlideTransform } from './correction/conte
 import { ContextTransition } from './correction/context-transition.js';
 import { ExecutionTimer } from './correction/execution-timer.js';
 import { ModelCompositor } from './model-compositor.js';
-import { getBestTokenMatches } from './correction/distance-modeler.js';
+import { EDIT_DISTANCE_COST_SCALE, getBestTokenMatches } from './correction/distance-modeler.js';
 import { TokenResultMapping } from './correction/token-result-mapping.js';
 
 import CasingForm = LexicalModelTypes.CasingForm;
@@ -245,11 +245,24 @@ export async function correctAndEnumerateWithoutTraversals(
     });
   }
 
-  // Remove `null` entries.
-  predictionRoots = predictionRoots.filter(tuple => !!tuple);
+  const wordbreak = determineModelWordbreaker(lexicalModel);
+  // Remove `null` entries, then determine suggestions.
+  predictionRoots.forEach((pr) => {
+    const postContext = models.applyTransform(pr.sample, context);
+    const tailTokenText = wordbreak(postContext);
+    const rootContext = models.applyTransform({insert: '', deleteLeft: KMWString.length(tailTokenText)}, postContext);
+
+    const results = predictFromCorrectionSequence(lexicalModel, [{
+      sample: {
+        insert: tailTokenText,
+        deleteLeft: 0,
+        id: pr.sample.id
+      },
+      p: pr.p
+    }], rootContext);
+    results.forEach((r) => rawPredictions.push(r));
+  })
 
-  // Running in bulk over all suggestions, duplicate entries may be possible.
-  rawPredictions = predictFromCorrections(lexicalModel, predictionRoots, context);
   if(allowSpace) {
     rawPredictions.forEach((entry) => entry.preservationTransform = inputTransform);
   }
@@ -449,13 +462,7 @@ export function buildAndMapPredictions(
   const { tokensToRemove, tokensToPredict } = determineSuggestionRange(applicationTarget.tokens, tokenization.tokens, (a, b) => a.spaceId == b.spaceId);
 
   const deleteLeft = tokensToPredict.length > 1 ? 0 : tokensToRemove.reduce((prev, curr) => prev + curr.searchModule.codepointLength, 0);
-
-  // Exists to be extended by the 'correctionTransfrom' below.
-  const emptyContext: Context = {
-    left: '',
-    startOfBuffer: false,
-    endOfBuffer: false
-  };
+  const rootContext = models.applyTransform({insert: '', deleteLeft}, transition.base.context);
 
   // Replace the existing context with the correction.
   const correctionTransform: Transform = {
@@ -470,7 +477,7 @@ export function buildAndMapPredictions(
     p: Math.exp(-rootCost * costFactor)
   };
 
-  let predictions = predictFromCorrections(model, [predictionRoot], emptyContext);
+  const predictions = predictFromCorrectionSequence(model, [predictionRoot], rootContext);
   predictions.forEach((entry) => {
     entry.preservationTransform = tokenization.taillessTrueKeystroke;
     // // Will need an extra lookup layer if the suggestion is generated from within a cluster.
@@ -664,43 +671,94 @@ export function shouldStopSearchingEarly(
  * @param context
  * @returns
  */
-export function predictFromCorrections(
+export function predictFromCorrectionSequence(
   lexicalModel: LexicalModel,
   corrections: ProbabilityMass<Transform>[],
   context: Context
 ): CorrectionPredictionTuple[] {
-  let returnedPredictions: CorrectionPredictionTuple[] = [];
-  const wordbreak = determineModelWordbreaker(lexicalModel);
+  let predictionPrefixSequence: ProbabilityMass<Suggestion>[] = [];
+  let tailPredictions: ProbabilityMass<Suggestion>[];
 
-  for(let correction of corrections) {
-    let predictions = lexicalModel.predict(correction.sample, context);
+  let currentContext = context;
+  let successfulPredictions = 0;
 
-    const { sample: correctionTransform, p: correctionProb } = correction;
-    const correctionRoot = wordbreak(models.applyTransform(correction.sample, context));
+  for(let i = 0; i < corrections.length; i++) {
+    const correction = corrections[i].sample;
 
-    let predictionSet = predictions.map((pair: ProbabilityMass<Suggestion>) => {
-      // Let's not rely on the model to copy transform IDs.
-      // Only bother is there IS an ID to copy.
-      if(correctionTransform.id !== undefined) {
-        pair.sample.transformId = correctionTransform.id;
+    // Step 2:  predict based on the final token.
+    const predictions = lexicalModel.predict(correction, currentContext);
+
+    // Failsafe:  if there are no matching predictions, create a fake prediction
+    // matching the original text.
+    if(predictions.length != 0) {
+      successfulPredictions++;
+    } else {
+      predictions.push({
+        sample: {
+          transform: correction,
+          displayAs: correction.insert
+        },
+        // It's not found in the lexicon, so we'll take a low probability for it.
+        //
+        // Edit penalties will be applied via the correction component separately later on.
+        p: -Math.exp(EDIT_DISTANCE_COST_SCALE)
+      });
+    }
+
+    if(i == corrections.length - 1) {
+      tailPredictions = predictions;
+    } else {
+      let bestMatch = predictions.find((p) => KMWString.length(p.sample.transform.insert) == KMWString.length(correction.insert));
+      if(!bestMatch) {
+        bestMatch = predictions[0];
       }
 
-      let tuple: CorrectionPredictionTuple = {
-        prediction: pair,
-        correction: {
-          sample: correctionRoot,
-          p: correctionProb
-        },
-        totalProb: pair.p * correctionProb,
-        matchLevel: SuggestionSimilarity.none
-      };
-      return tuple;
-    });
+      predictionPrefixSequence = predictionPrefixSequence.concat(bestMatch);
+    }
 
-    returnedPredictions = returnedPredictions.concat(predictionSet);
+    // Or maybe per prediction, in some manner?
+    currentContext = models.applyTransform(correction, currentContext);
   }
 
-  return returnedPredictions;
+  if(!successfulPredictions) {
+    return [];
+  }
+
+  const predictions: CorrectionPredictionTuple[] = tailPredictions.map((p) => {
+    // Concat corrections + predictions for their components.
+    const predictionSequence = [...predictionPrefixSequence, p];
+    const fullPrediction: ProbabilityMass<Suggestion> = predictionSequence.reduce((prev, curr) => {
+      return {
+        sample: {
+          transform: models.buildMergedTransform(prev.sample.transform, curr.sample.transform),
+          displayAs: prev.sample.displayAs + curr.sample.displayAs
+        },
+        p: prev.p * curr.p
+      };
+    }, {sample: {transform: {insert: '', deleteLeft: 0}, displayAs: ''}, p: 1});
+
+    const fullCorrection: ProbabilityMass<string> = corrections.reduce((prev, curr) => {
+      return {
+        sample: prev.sample + curr.sample.insert,
+        p: prev.p * curr.p
+      }
+    }, {sample: '', p: 1})
+
+    const transformId = p.sample.transform.id;
+    if(transformId) {
+      fullPrediction.sample.transform.id = transformId;
+      fullPrediction.sample.transformId = transformId;
+    }
+
+    return {
+      prediction: fullPrediction,
+      correction: fullCorrection,
+      totalProb: fullPrediction.p * fullCorrection.p,
+      matchLevel: SuggestionSimilarity.none,
+    };
+  });
+
+  return predictions;
 }
 
 /**

web/src/test/auto/headless/engine/predictive-text/worker-thread/prediction-helpers/predict-from-corrections.tests.ts

@@ -4,7 +4,7 @@ import { assert } from 'chai';
 import { deepCopy } from "keyman/common/web-utils";
 import { LexicalModelTypes } from '@keymanapp/common-types';
 
-import { models, predictFromCorrections, tupleDisplayOrderSort } from "@keymanapp/lm-worker/test-index";
+import { models, predictFromCorrectionSequence, tupleDisplayOrderSort } from "@keymanapp/lm-worker/test-index";
 
 import CasingFunction = LexicalModelTypes.CasingFunction;
 import Context = LexicalModelTypes.Context;
@@ -71,10 +71,10 @@ const DUMMY_MODEL_CONFIG = {
   languageUsesCasing: true
 };
 
-describe('predictFromCorrections', () => {
+describe('predictFromCorrectionSequence', () => {
   it('handles a single correction prefixing multiple entries - no transform ID', () => {
     const context: Context = {
-      left: 'It',
+      left: '',
       right: '',
       startOfBuffer: true,
       endOfBuffer: true
@@ -82,7 +82,7 @@ describe('predictFromCorrections', () => {
 
     const correctionDistribution: Distribution<Transform> = [{
         sample: {
-          insert: 's',
+          insert: 'Its',
           deleteLeft: 0
         },
         p: 0.6
@@ -112,12 +112,15 @@ describe('predictFromCorrections', () => {
       futureSuggestions: [ dummied_suggestions ]
     });
 
-    const predictions = predictFromCorrections(model, correctionDistribution, context);
+    const predictions = predictFromCorrectionSequence(model, correctionDistribution, context);
     predictions.forEach((entry) => assert.equal(entry.correction.sample, 'Its'));
     predictions.forEach((entry) => assert.equal(entry.correction.p, 0.6));
     predictions.sort(tupleDisplayOrderSort);
 
-    assert.sameDeepOrderedMembers(predictions.map((entry) => entry.prediction.sample), dummied_suggestions);
+    assert.sameDeepOrderedMembers(predictions.map((entry) => entry.prediction.sample), dummied_suggestions.map((s) => {
+      delete s.p;
+      return s;
+    }));
 
     assert.approximately(predictions[0].totalProb, 0.18 * 0.6, 0.00001);
     assert.approximately(predictions[1].totalProb, 0.02 * 0.6, 0.00001);
@@ -125,7 +128,7 @@ describe('predictFromCorrections', () => {
 
   it('handles a single correction prefixing multiple entries - with transform ID', () => {
     const context: Context = {
-      left: 'It',
+      left: '',
       right: '',
       startOfBuffer: true,
       endOfBuffer: true
@@ -133,7 +136,7 @@ describe('predictFromCorrections', () => {
 
     const correctionDistribution: Distribution<Transform> = [{
         sample: {
-          insert: 's',
+          insert: 'Its',
           deleteLeft: 0,
           id: 314159
         },
@@ -164,7 +167,7 @@ describe('predictFromCorrections', () => {
       futureSuggestions: [ dummied_suggestions ]
     });
 
-    const predictions = predictFromCorrections(model, correctionDistribution, context);
+    const predictions = predictFromCorrectionSequence(model, correctionDistribution, context);
     predictions.forEach((entry) => assert.equal(entry.correction.sample, 'Its'));
     predictions.forEach((entry) => assert.equal(entry.correction.p, 0.6));
     predictions.sort(tupleDisplayOrderSort);
@@ -173,89 +176,11 @@ describe('predictFromCorrections', () => {
     assert.sameDeepOrderedMembers(predictions.map((entry) => entry.prediction.sample), dummied_suggestions.map((entry) => {
       entry = deepCopy(entry);
       entry.transformId = 314159;
+      delete entry.p;
       return entry;
     }));
 
     assert.approximately(predictions[0].totalProb, 0.18 * 0.6, 0.00001);
     assert.approximately(predictions[1].totalProb, 0.02 * 0.6, 0.00001);
   });
-
-  it('handles multiple corrections at once', () => {
-    const context: Context = {
-      left: 'It',
-      right: '',
-      startOfBuffer: true,
-      endOfBuffer: true
-    };
-
-    // Note:  each correction is used in order in a separate model.predict call.
-    /** @type {Distribution<Transform>} */
-    const correctionDistribution: Distribution<Transform> = [{
-        // postContext:  is
-        sample: {
-          insert: 's',
-          deleteLeft: 1
-        },
-        p: 0.4
-      }, {
-        // postContext: its
-        sample: {
-          insert: 's',
-          deleteLeft: 0
-        },
-        p: 0.6
-      }
-    ];
-
-    const dummied_suggestions: Outcome<Suggestion>[][] = [
-      // postContext:  is
-      [{
-        transform: {
-          insert: "is",
-          deleteLeft: 2
-        },
-        displayAs: "is",
-        p: 0.4
-      }, {
-        transform: {
-          insert: "isn't",
-          deleteLeft: 2
-        },
-        displayAs: "isn't",
-        p: 0.2
-      }],
-      // postContext: its
-      [{
-        transform: {
-          insert: "it's",
-          deleteLeft: 2
-        },
-        displayAs: "it's",
-        p: 0.18
-      }, {
-        transform: {
-          insert: "its",
-          deleteLeft: 2
-        },
-        displayAs: "its",
-        p: 0.02
-      }]
-    ];
-
-    const model = new DummyModel({
-      ...DUMMY_MODEL_CONFIG,
-      futureSuggestions: dummied_suggestions
-    });
-
-    const predictions = predictFromCorrections(model, correctionDistribution, context);
-    predictions.sort(tupleDisplayOrderSort);
-
-    assert.sameOrderedMembers(predictions.map((entry) => entry.prediction.sample.displayAs), ["is", "it's", "isn't", "its"]);
-    assert.sameDeepMembers(predictions.map((entry) => entry.prediction.sample), dummied_suggestions.flatMap((entry) => entry));
-
-    assert.approximately(predictions[0].totalProb, 0.4 * 0.4, 0.00001);
-    assert.approximately(predictions[1].totalProb, 0.18 * 0.6, 0.00001);
-    assert.approximately(predictions[2].totalProb, 0.4 * 0.2, 0.00001);
-    assert.approximately(predictions[3].totalProb, 0.02 * 0.6, 0.00001);
-  });
 });

web/src/test/auto/headless/engine/predictive-text/worker-thread/worker-custom-punctuation.tests.ts

@@ -81,6 +81,24 @@ describe('Custom Punctuation', function () {
             open: "'",
             close: "'"
           }
+        },
+        // Some of the suggestions above actually wordbreak differently from
+        // what might be expected.  So, we override the wordbreaker to ensure
+        // the tests run smoothly.
+        wordbreaker: (text) => {
+          const textLen = text.length;
+          if(text.charAt(textLen - 1) == " ") {
+            return [
+              {text: text.substring(0, 1), start: 0, end: 1, length: 1},
+              {text: text.substring(1, textLen-2), start: 1, end: textLen-1, length: textLen-2},
+              {text: text.substring(textLen-1), start: textLen-1, end: textLen, length: 1}
+            ];
+          } else {
+            return [
+              {text: text.substring(0, 1), start: 0, end: 1, length: 1},
+              {text: text.substring(1), start: 1, end: textLen, length: textLen-1}
+            ];
+          }
         }
       });