gpt-token-utils/dist/EncoderResult.mjs

Isomorphic utilities for GPT-3 tokenization and prompt building.

/**
 * @copyright Sister Software. All rights reserved.
 * @author Teffen Ellis, et al.
 * @license
 * See LICENSE file in the project root for full license information.
 */
const nodeInspectSymbol = Symbol.for('nodejs.util.inspect.custom');
const supportsSegmenter = typeof Intl !== 'undefined' && typeof Intl.Segmenter !== 'undefined';
/**
 * The `EncoderResult` includes information for post-encoding analysis such as...
 *
 * - The tokens that were encoded.
 * - The BPE token pairs that were used during encoded.
 * - Two-way maps of tokens to BPE token pairs.
 *
 * This information can be used to analyze the encoding process and to
 * reconstruct the original string from the encoded tokens.
 *
 * Note that this object is considered immutable. Consider encoding a new string
 * if you need an updated `EncoderResult`.
 *
 * @see {@linkcode BytePairEncoder}
 */
export class EncoderResult {
    /**
     * A map of BPE token pairs to the corresponding token.
     */
    tokenBPEMap;
    /**
     * A map of tokens to the corresponding BPE token pair.
     */
    bpeTokenMap;
    /**
     * A map of BPE token pairs to the number of times they were used during encoding.
     * The key is the BPE token pair and the value is the number of times it appeared.
     */
    bpeCountsMap;
    /**
     * A map of tokens to the number of times they were used during encoding.
     * The key is the token and the value is the number of times it appeared.
     */
    tokenCountsMap;
    tokens;
    bpeTokenPairs;
    originalInput;
    matchedTextSegments;
    segmenter;
    constructor({ tokens, bpeTokenPairs, originalInput, matchedTextSegments }, locale) {
        if (bpeTokenPairs.length !== tokens.length) {
            throw new Error('The number of BPE token pairs must match the number of tokens.');
        }
        const tokenToBPE = [];
        const BPEToToken = [];
        const tokenCountsMap = new Map();
        const bpeCountsMap = new Map();
        for (let i = 0; i < bpeTokenPairs.length; i++) {
            const token = tokens[i];
            const bpeTokenPair = bpeTokenPairs[i];
            const tokenCount = tokenCountsMap.get(token) || 0;
            const bpeCount = bpeCountsMap.get(bpeTokenPair) || 0;
            tokenCountsMap.set(token, tokenCount + 1);
            bpeCountsMap.set(bpeTokenPair, bpeCount + 1);
            tokenToBPE.push([token, bpeTokenPair]);
            BPEToToken.push([bpeTokenPair, token]);
        }
        this.tokenBPEMap = new Map(tokenToBPE);
        this.bpeTokenMap = new Map(BPEToToken);
        this.tokenCountsMap = tokenCountsMap;
        this.bpeCountsMap = bpeCountsMap;
        this.tokens = tokens;
        this.bpeTokenPairs = bpeTokenPairs;
        this.originalInput = originalInput;
        this.matchedTextSegments = matchedTextSegments;
        if (supportsSegmenter) {
            this.segmenter = new Intl.Segmenter(locale);
        }
    }
    /**
     * Get the encoded byte-pair for a given token.
     */
    getBPE(token) {
        return this.tokenBPEMap.get(token);
    }
    /**
     * Get the number of times a given token appeared during encoding.
     * @see {@linkcode EncoderResult.length} if you're just trying count number of tokens.
     */
    getTokenCount(token) {
        return this.tokenCountsMap.get(token) || 0;
    }
    /**
     * Get the number of times a given byte-pair appeared during encoding.
     */
    getBPECount(bpe) {
        return this.bpeCountsMap.get(bpe) || 0;
    }
    /**
     * Iterate over the tokens in the result.
     */
    [Symbol.iterator]() {
        return this.tokens[Symbol.iterator]();
    }
    /**
     * The number of tokens in the result.
     */
    get length() {
        return this.tokens.length;
    }
    /**
     * The number of characters in the original text.
     *
     * @see {@link https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Intl/Segmenter Intl.Segmenter}
     */
    get characterCount() {
        if (!this.segmenter) {
            console.warn('Intl.Segmenter is not supported. Falling back to string length.');
            return this.originalInput.length;
        }
        return Array.from(this.segmenter.segment(this.originalInput)).length;
    }
    [nodeInspectSymbol]() {
        return `EncoderResult(${this.length})`;
    }
    toString() {
        return this[Symbol.iterator]().toString();
    }
    toJSON() {
        return {
            tokens: this.tokens,
            bpeTokenPairs: this.bpeTokenPairs,
            originalInput: this.originalInput,
            matchedTextSegments: this.matchedTextSegments,
        };
    }
}