File size: 2,348 Bytes
74e8f2f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
/**
* @license
* Copyright Big Vision Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Copied from
// https://github.com/tensorflow/tfjs-models/blob/master/universal-sentence-encoder/src/tokenizer/trie.ts
import {stringToChars} from './common';
// [token, score, index]
type OutputNode = [string[], number, number];
class TrieNode {
parent: TrieNode|null;
end: boolean;
children: {[firstSymbol: string]: TrieNode};
word: OutputNode;
constructor() {
this.parent = null;
this.children = {};
this.end = false;
this.word = [[], 0, 0];
}
}
/**
* Simple Trie datastructure.
*/
export class Trie {
root: TrieNode;
constructor() {
this.root = new TrieNode();
}
/**
* Inserts a token into the trie.
*/
insert(word: string, score: number, index: number) {
let node = this.root;
const symbols = stringToChars(word);
for (let i = 0; i < symbols.length; i++) {
if (!node.children[symbols[i]]) {
node.children[symbols[i]] = new TrieNode();
node.children[symbols[i]].parent = node;
node.children[symbols[i]].word[0] = node.word[0].concat(symbols[i]);
}
node = node.children[symbols[i]];
if (i === symbols.length - 1) {
node.end = true;
node.word[1] = score;
node.word[2] = index;
}
}
}
/**
* Returns an array of all tokens starting with ss.
*
* @param ss The prefix to match on.
*/
commonPrefixSearch(ss: string[]): OutputNode[] {
const output: OutputNode[] = [];
let node = this.root.children[ss[0]];
for (let i = 0; i < ss.length && node; i++) {
if (node.end) {
output.push(node.word);
}
node = node.children[ss[i + 1]];
}
if (!output.length) {
output.push([[ss[0]], 0, 0]);
}
return output;
}
}
|