gpt-token-utils
Version:
Isomorphic utilities for GPT-3 tokenization and prompt building.
78 lines (77 loc) • 2.56 kB
JavaScript
/**
* @copyright Sister Software. All rights reserved.
* @author Teffen Ellis, et al.
* @license
* See LICENSE file in the project root for full license information.
*/
/**
* Parses a BPE file into a list of bigrams
*
* The vocab.bpe file is a text file that contains a set of byte pair encoding (BPE) codes
* that are used in the tokenization process.
*
* The file should be in the following format:
*
* ```text
* #version: VERSION_STRING
* [prefix1] [suffix1]
* [prefixN] [suffixN]
* ...
* ```
*/
export function parseBPEFile(bpeFileContents) {
const lines = bpeFileContents.trim().split('\n');
const [versionLine, ...bpeMerges] = lines;
const [version = 'unknown'] = versionLine.trim().match(/^#version: (\d.+)$/) || [];
const entries = bpeMerges.map((line, lineIndex) => {
const segments = line
// Each line contains a pair of tokens separated by a space
.split(/(\s+)/)
// Clean up the tokens...
.map((x) => x.trim())
.filter(Boolean);
if (segments.length < 2) {
throw new Error(`Invalid BPE file format: line ${lineIndex + 1} is not a valid bigram`);
}
const [prefix, suffix] = segments;
const entry = {
prefix,
suffix,
};
return entry;
});
return {
version,
entries,
};
}
/**
* Parse a token encoder file, usually from a file named `encoder.json`
*/
export function parseEncoderFile(
/**
* The token encoder content, either as a string or as a parsed object.
*/
tokenEncoderContent) {
const tokenEncodings = typeof tokenEncoderContent === 'string' ? JSON.parse(tokenEncoderContent) : tokenEncoderContent;
return tokenEncodings;
}
export function parseTikTokenFile(tikTokenFileContents) {
const lines = tikTokenFileContents.trim().split('\n');
const entries = lines.map((line, lineIndex) => {
const segments = line
// Each line contains a pair of tokens separated by a space
.split(/(\s+)/)
// Clean up the tokens...
.map((x) => x.trim())
.filter(Boolean);
if (segments.length < 2) {
throw new Error(`Invalid tiktoken file format: line ${lineIndex + 1} is not a valid bigram`);
}
const [encodedToken, rankString] = segments;
const token = Buffer.from(encodedToken, 'base64').toString('utf8');
const rank = parseInt(rankString, 10);
return { token, rank };
});
return entries;
}