diff --git a/web/src/engine/predictive-text/worker-thread/src/main/predict-helpers.ts b/web/src/engine/predictive-text/worker-thread/src/main/predict-helpers.ts index b8bf4b1570..bd06a89e70 100644 --- a/web/src/engine/predictive-text/worker-thread/src/main/predict-helpers.ts +++ b/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[], context: Context ): CorrectionPredictionTuple[] { - let returnedPredictions: CorrectionPredictionTuple[] = []; - const wordbreak = determineModelWordbreaker(lexicalModel); + let predictionPrefixSequence: ProbabilityMass[] = []; + let tailPredictions: ProbabilityMass[]; - 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) => { - // 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 = 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 = 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; } /** diff --git a/web/src/test/auto/headless/engine/predictive-text/worker-thread/prediction-helpers/predict-from-corrections.tests.ts b/web/src/test/auto/headless/engine/predictive-text/worker-thread/prediction-helpers/predict-from-corrections.tests.ts index 2351ff9932..35f69975be 100644 --- a/web/src/test/auto/headless/engine/predictive-text/worker-thread/prediction-helpers/predict-from-corrections.tests.ts +++ b/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 = [{ 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 = [{ 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} */ - const correctionDistribution: Distribution = [{ - // postContext: is - sample: { - insert: 's', - deleteLeft: 1 - }, - p: 0.4 - }, { - // postContext: its - sample: { - insert: 's', - deleteLeft: 0 - }, - p: 0.6 - } - ]; - - const dummied_suggestions: Outcome[][] = [ - // 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); - }); }); \ No newline at end of file diff --git a/web/src/test/auto/headless/engine/predictive-text/worker-thread/worker-custom-punctuation.tests.ts b/web/src/test/auto/headless/engine/predictive-text/worker-thread/worker-custom-punctuation.tests.ts index 9b9ab2c312..cd4dbd106e 100644 --- a/web/src/test/auto/headless/engine/predictive-text/worker-thread/worker-custom-punctuation.tests.ts +++ b/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} + ]; + } } });