Merge pull request #11870 from keymanapp/change/common/models/templates/trie-results-through-traversal

change(common/models/templates): rework Trie predict method to utilize traversals 📚
This commit is contained in:
Joshua Horton 2024-07-08 14:37:36 +07:00 • committed by GitHub
commit 479ea32dd0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 46 additions and 116 deletions

View file

@ -73,15 +73,6 @@ export interface TrieModelOptions {
punctuation?: LexicalModelPunctuation;
}
/**
* Used to determine the probability of an entry from the trie.
*/
type TextWithProbability = {
text: string;
// TODO: use negative-log scaling instead?
p: number; // real-number weight, from 0 to 1
}
class Traversal implements LexiconTraversal {
/**
* The lexical prefix corresponding to the current traversal state.
@ -378,11 +369,10 @@ export default class TrieModel implements LexicalModel {
type SearchKey = string & { _: 'SearchKey'};
/**
* The priority queue will always pop the most weighted item. There can only
* be two kinds of items right now: nodes, and entries; both having a weight
* attribute.
* The priority queue will always pop the most probable item - be it a Traversal
* state or a lexical entry reached via Traversal.
*/
type Weighted = Node | Entry;
type TraversableWithProb = TextWithProbability | LexiconTraversal;
/**
* A function that converts a string (word form or query) into a search key
@ -465,12 +455,8 @@ class Trie {
*/
lookup(prefix: string): TextWithProbability[] {
let searchKey = this.toKey(prefix);
let lowestCommonNode = findPrefix(this.root, searchKey);
if (lowestCommonNode === null) {
return [];
}
return getSortedResults(lowestCommonNode, searchKey, this.totalWeight);
let rootTraversal = this.traverseFromRoot().child(searchKey);
return rootTraversal ? getSortedResults(rootTraversal) : [];
}
/**
@ -478,36 +464,10 @@ class Trie {
* @param n How many suggestions, maximum, to return.
*/
firstN(n: number): TextWithProbability[] {
return getSortedResults(this.root, '' as SearchKey, this.totalWeight, n);
return getSortedResults(this.traverseFromRoot(), n);
}
}
/**
* Finds the deepest descendent in the trie with the given prefix key.
*
* This means that a search in the trie for a given prefix has a best-case
* complexity of O(m) where m is the length of the prefix.
*
* @param key The prefix to search for.
* @param index The index in the prefix. Initially 0.
*/
function findPrefix(node: Node, key: SearchKey, index: number = 0): Node | null {
// An important note - the Trie itself is built on a per-JS-character basis,
// not on a UTF-8 character-code basis.
if (node.type === 'leaf' || index === key.length) {
return node;
}
// So, for SMP models, we need to match each char of the supplementary pair
// in sequence. Each has its own node in the Trie.
let char = key[index];
if (node.children[char]) {
return findPrefix(node.children[char], key, index + 1);
}
return null;
}
/**
* Returns all entries matching the given prefix, in descending order of
* weight.
@ -516,72 +476,36 @@ function findPrefix(node: Node, key: SearchKey, index: number = 0): Node | null
* @param results the current results
* @param queue
*/
function getSortedResults(node: Node, prefix: SearchKey, N: number, limit = MAX_SUGGESTIONS): TextWithProbability[] {
let queue = new PriorityQueue(function(a: Weighted, b: Weighted) {
function getSortedResults(traversal: LexiconTraversal, limit = MAX_SUGGESTIONS): TextWithProbability[] {
let queue = new PriorityQueue(function(a: TraversableWithProb, b: TraversableWithProb) {
// In case of Trie compilation issues that emit `null` or `undefined`
return (b ? b.weight : 0) - (a ? a.weight : 0);
return (b ? b.p : 0) - (a ? a.p : 0);
});
let results: TextWithProbability[] = [];
if (node.type === 'leaf') {
// Assuming the values are sorted, we can just add all of the values in the
// leaf, until we reach the limit.
for (let item of node.entries) {
// String.startsWith is not supported on certain Android (5.0) devices we wish to support.
// Requires a minimum of Chrome 36, as opposed to 5.0's default of 35.
if (item.key.indexOf(prefix) == 0) {
let { content, weight } = item;
results.push({
text: content,
p: weight / N
});
queue.enqueue(traversal);
if (results.length >= limit) {
return results;
}
}
}
} else {
queue.enqueue(node);
let next: Weighted | undefined;
while(queue.count > 0) {
const entry = queue.dequeue();
while (next = queue.dequeue()) {
if (isNode(next)) {
// When a node is next up in the queue, that means that next least
// likely suggestion is among its decsendants.
// So we search all of its descendants!
if (next.type === 'leaf') {
queue.enqueueAll(next.entries);
} else {
// XXX: alias `next` so that TypeScript can be SURE that internal is
// in fact an internal node. Because of the callback binding to the
// original definition of node (i.e., a Node | Entry), this will not
// type-check otherwise.
let internal = next;
queue.enqueueAll(next.values.map(char => {
return internal.children[char];
}));
}
} else {
// When an entry is up next in the queue, we just add its contents to
// the results!
results.push({
text: next.content,
p: next.weight / N
});
if (results.length >= limit) {
return results;
}
if((entry as TextWithProbability)!.text !== undefined) {
const lexicalEntry = entry as TextWithProbability;
results.push(lexicalEntry);
if(results.length >= limit) {
return results;
}
} else {
const traversal = entry as LexiconTraversal;
queue.enqueueAll(traversal.entries);
let children: LexiconTraversal[] = []
for(let child of traversal.children()) {
children.push(child.traversal());
}
queue.enqueueAll(children);
}
}
return results;
}
/** TypeScript type guard that returns whether the thing is a Node. */
function isNode(x: Entry | Node): x is Node {
return 'type' in x;
}
/**

View file

@ -19,6 +19,25 @@ declare type USVString = string;
declare type CasingForm = 'lower' | 'initial' | 'upper';
/**
* Represents one lexical entry and its probability..
*/
type TextWithProbability = {
/**
* A lexical entry (word) offered by the model.
*
* Note: not the search-term keyed part. This will match the actual, unkeyed form.
*/
text: string;
/**
* The probability of the lexical entry, directly based upon its frequency.
*
* A real-number weight, from 0 to 1.
*/
p: number;
}
/**
* Used to facilitate edit-distance calculations by allowing the LMLayer to
* efficiently search the model's lexicon in a Trie-like manner.
@ -56,8 +75,6 @@ declare interface LexiconTraversal {
* This allows bypassing iteration among all legal child Traversals.
*
* If such a traversal state is not supported, returns `undefined`.
* Implementations may choose to return `undefined` if more than one UTF-16
* codepoint is appended, even if such a descendant exists.
*
* Note: traversals navigate and represent the lexicon in its "keyed" state,
* as produced by use of the search-term keying function defined for the model.
@ -86,18 +103,7 @@ declare interface LexiconTraversal {
* - prefix of 'crepe': ['crêpe', 'crêpé']
* - other examples: https://www.thoughtco.com/french-accent-homographs-1371072
*/
entries: {
/**
* A lexical entry (word) offered by the model.
*
* Note: not the search-term keyed part. This will match the actual, unkeyed form.
*/
text: USVString,
/**
* The probability of the lexical entry, directly based upon its frequency.
*/
p: number
}[];
entries: TextWithProbability[];
// Note: `p`, not `maxP` - we want to see the same name for `this.entries.p` and `this.p`
/**