markov-coil/dist/MarkovCoil.js

Markov chain optimized for large bodies of text.

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.MarkovCoil = exports.MarkovNode = void 0;
class MarkovNode {
    constructor() {
        this.children = new Map();
        this.weight = 0;
    }
}
exports.MarkovNode = MarkovNode;
class MarkovCoil {
    /**
     * Create new markov chain for token prediction.
     * @param {string[]} tokens Tokens to generate chain with.
     * @param {number} [depth=3] Depth of the chain. A depth of n will use n tokens in its search for a prediction (n-gram).
     */
    constructor(tokens, depth = 3) {
        this.vocab = {
            tokens: [],
            indexes: new Map(),
        };
        this.root = new MarkovNode();
        this.depth = depth;
        this.createChain(tokens);
    }
    getIndex(token) {
        return this.vocab.indexes.get(token);
    }
    getToken(index) {
        return this.vocab.tokens[index];
    }
    addTokenToVocab(token) {
        const index = this.getIndex(token);
        if (index === undefined) {
            this.vocab.tokens.push(token);
            this.vocab.indexes.set(token, this.vocab.tokens.length - 1);
        }
    }
    addSequenceToChain(sequence) {
        let current = this.root;
        sequence.forEach((token) => {
            current.weight += 1;
            const index = this.getIndex(token);
            if (current.children.has(index) === false) {
                current.children.set(index, new MarkovNode());
            }
            current = current.children.get(index);
        });
        current.weight += 1;
    }
    createChain(tokens) {
        tokens.forEach((token) => {
            this.addTokenToVocab(token);
        });
        for (let i = 0; i < tokens.length; i += 1) {
            const sequence = tokens.slice(i, i + this.depth + 1);
            this.addSequenceToChain(sequence);
        }
    }
    /**
     * Predict next possible values for the given sequence of tokens.
     * @param {string[]} context An array of tokens.
     * @returns A mapping of possible tokens to their probability of occuring.
     */
    predictions(context) {
        const indexes = context
            .slice(-1 * this.depth)
            .map((token) => this.getIndex(token));
        let current = this.root;
        indexes.forEach((index) => {
            if (index === undefined || current.children.has(index) === false) {
                current = this.root;
                return;
            }
            current = current.children.get(index);
        });
        const result = {};
        current.children.forEach((node, index) => {
            result[this.getToken(index)] = node.weight / current.weight;
        });
        return result;
    }
    weightedChoice(predictions) {
        const tokens = Object.keys(predictions);
        if (tokens.length === 0)
            return null;
        // Accumulate weights
        const weights = tokens.map((token) => predictions[token]);
        for (let i = 1; i < weights.length; i += 1) {
            weights[i] += weights[i - 1];
        }
        const random = Math.random() * weights[weights.length - 1];
        // Could probably optimize this better with binary search, since weights are now ordered
        for (let i = 0; i < tokens.length; i += 1) {
            if (weights[i] > random) {
                return tokens[i];
            }
        }
        // Fallback (should not reach)
        return null;
    }
    /**
     * Predict the next token given a starting sequence.
     * @param {string[]} sequence The starting sequence of tokens.
     * @param {boolean} [weighted=true] If true, will use weighted random choice from all possible predictions. If false, will use random choice.
     * @returns {string} The predicted token.
     */
    predict(sequence, weighted = true) {
        const predictions = this.predictions(sequence);
        const tokens = Object.keys(predictions);
        if (tokens.length === 0)
            return null;
        if (!weighted)
            return tokens[Math.floor(Math.random() * tokens.length)];
        return this.weightedChoice(predictions);
    }
    /**
     * Predicts a sequence of tokens that could follow the starting sequence.
     * @param {token[]} sequence The starting sequence of tokens.
     * @param {number} length The length of the predicted sequence.
     * @param {boolean} [weighted=true] If true, will use weighted random choice from all possible predictions. If false, will use random choice.
     * @returns {string[]} A sequence of predicted tokens.
     */
    predictSequence(sequence, length, weighted = true) {
        if (length === 0) {
            return [];
        }
        const prediction = this.predict(sequence, weighted);
        if (prediction === null) {
            return this.predictSequence([], length - 1, weighted);
        }
        const nextSequence = sequence.slice(1).concat(prediction);
        return [prediction].concat(this.predictSequence(nextSequence, length - 1, weighted));
    }
    /** Prints tabulated trie structure to console. Useful for debugging. */
    prettyPrint(useVocab = false) {
        console.log(`Root (${this.root.weight})`);
        const prettyPrintHelper = (node, spaces = 2) => {
            node.children.forEach((child, index) => {
                let line = " ".repeat(spaces);
                if (useVocab) {
                    line += this.getToken(index);
                }
                else {
                    line += index;
                }
                line += ` (${child.weight})`;
                console.log(line);
                prettyPrintHelper(child, spaces + 2);
            });
        };
        prettyPrintHelper(this.root);
    }
}
exports.MarkovCoil = MarkovCoil;