gpt-token-utils/dist/RanksMap.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');
/**
 * Map of byte-pair encodings according to their BPE rank
 * @internal
 */
export class RanksMap {
    _prefixToSuffixRankMap = new Map();
    getRank(prefix, suffix) {
        const suffixMap = this._prefixToSuffixRankMap.get(prefix);
        if (suffixMap) {
            return suffixMap.get(suffix);
        }
    }
    constructor(vocab, mergesSpacesCount = 0) {
        const normalizedVocab = Array.isArray(vocab) ? vocab.slice() : vocab.entries.slice();
        if (mergesSpacesCount > 0) {
            for (let i = 1; i < mergesSpacesCount; i++) {
                for (let j = 1; j < mergesSpacesCount; j++) {
                    if (i + j <= mergesSpacesCount) {
                        normalizedVocab.push({
                            prefix: '\u0120'.repeat(i),
                            suffix: '\u0120'.repeat(j),
                        });
                    }
                }
            }
        }
        for (const [rank, entry] of normalizedVocab.entries()) {
            let suffixMap = this._prefixToSuffixRankMap.get(entry.prefix);
            if (!suffixMap) {
                suffixMap = new Map();
                this._prefixToSuffixRankMap.set(entry.prefix, suffixMap);
            }
            suffixMap.set(entry.suffix, rank);
        }
    }
    get size() {
        return this._prefixToSuffixRankMap.size;
    }
    [nodeInspectSymbol]() {
        return `RanksMap(${this.size})`;
    }
}