tokenx
Version:
Fast token estimation with 95%+ average accuracy, 2kB minified and gzipped
203 lines (202 loc) • 8.69 kB
JavaScript
// src/segments.ts
const PATTERNS = {
whitespace: /^\s+$/,
structuredWhitespace: /\n\s/,
nonAscii: /[\u0080-\uFFFF]/,
cjk: /[\u4E00-\u9FFF\u3400-\u4DBF\u3000-\u30FF\uFF00-\uFFEF\u2E80-\u2EFF\u31C0-\u31EF\u3200-\u32FF\u3300-\u33FF\uAC00-\uD7AF\u1100-\u11FF\u3130-\u318F\uA960-\uA97F\uD7B0-\uD7FF]/,
numeric: /^\d+$/,
punctuation: /[.,!?;(){}[\]<>:/\\|@#$%^&*+=`~_"-]/,
cjkPunctuation: /[\u3001-\u3003\u3008-\u3011\u3014-\u301F\u3030\u30FB\uFF01-\uFF0F\uFF1A-\uFF20\uFF3B-\uFF40\uFF5B-\uFF65]/,
lowercaseWord: /^[a-z]+$/
};
const TOKEN_SPLIT_PATTERN = new RegExp(`(\\s+|${PATTERNS.punctuation.source}+|${PATTERNS.cjkPunctuation.source}+)`);
const DEFAULT_CHARS_PER_TOKEN = 7;
const SHORT_TOKEN_THRESHOLD = 3;
const LOWERCASE_WORD_SINGLE_TOKEN_LENGTH = 8;
const PUNCTUATION_CHARS_PER_TOKEN = 6;
const KANA_CHARS_PER_TOKEN = 1.35;
const HANGUL_CHARS_PER_TOKEN = 1.65;
const HANZI_CHARS_PER_TOKEN = 1.2;
const DEFAULT_LANGUAGE_CONFIGS = [
{
pattern: /[äöüßẞ]/i,
averageCharsPerToken: 3
},
{
pattern: /[éèêëàâîïôûùüÿçœæáíóúñ]/i,
averageCharsPerToken: 4.5
},
{
pattern: /[ąćęłńóśźżěščřžýůúďťň]/i,
averageCharsPerToken: 2.5
},
{
pattern: /[\u0430-\u044F\u0451]/i,
averageCharsPerToken: 6
},
{
pattern: /[\u03AC-\u03CE]/i,
averageCharsPerToken: 3
},
{
pattern: /^\p{Extended_Pictographic}[\p{Extended_Pictographic}\p{Emoji_Component}]*$/u,
averageCharsPerToken: .9
}
];
/**
* Walks a text as (segment, estimated token count) pairs. The segments
* concatenate back to the original text; whitespace segments count zero
* tokens unless they carry structure (indentation or blank lines).
*/
function* walkSegments(text, options = {}) {
if (!text) return;
const resolvedOptions = resolveOptions(options);
let previousSegment = "";
const parts = text.split(TOKEN_SPLIT_PATTERN);
for (let i = 0; i < parts.length; i++) {
let segment = parts[i];
if (PATTERNS.cjkPunctuation.test(parts[i + 1] ?? "") && PATTERNS.cjk.test(segment)) segment += parts[++i];
if (segment) {
yield {
segment,
tokenCount: estimateSegmentTokens(segment, resolvedOptions, previousSegment)
};
previousSegment = segment;
}
}
}
function resolveOptions(options) {
return {
defaultCharsPerToken: options.defaultCharsPerToken ?? DEFAULT_CHARS_PER_TOKEN,
customLanguageConfigs: options.languageConfigs ?? []
};
}
function estimateSegmentTokens(segment, { customLanguageConfigs, defaultCharsPerToken }, previousSegment) {
if (PATTERNS.whitespace.test(segment)) {
if (PATTERNS.structuredWhitespace.test(segment)) return 1;
if (!segment.includes("\n")) return 0;
const previousCharacter = previousSegment.slice(-1);
return PATTERNS.punctuation.test(previousCharacter) || PATTERNS.cjkPunctuation.test(previousCharacter) ? 0 : 1;
}
const customCharsPerToken = findLanguageCharsPerToken(segment, customLanguageConfigs);
if (customCharsPerToken !== void 0) return Math.ceil(getCharacterCount(segment) / customCharsPerToken);
if (PATTERNS.cjk.test(segment)) return estimateCjkTokens(segment);
const builtInCharsPerToken = PATTERNS.nonAscii.test(segment) ? findLanguageCharsPerToken(segment, DEFAULT_LANGUAGE_CONFIGS) : void 0;
if (builtInCharsPerToken !== void 0) return Math.ceil(getCharacterCount(segment) / builtInCharsPerToken);
if (PATTERNS.numeric.test(segment)) return Math.ceil(segment.length / 3);
if (segment.length <= SHORT_TOKEN_THRESHOLD) return 1;
if (segment.length <= LOWERCASE_WORD_SINGLE_TOKEN_LENGTH && PATTERNS.lowercaseWord.test(segment)) return 1;
if (PATTERNS.punctuation.test(segment)) return Math.ceil(segment.length / PUNCTUATION_CHARS_PER_TOKEN);
return Math.ceil(segment.length / defaultCharsPerToken);
}
function findLanguageCharsPerToken(segment, languageConfigs) {
for (const config of languageConfigs) if (segment.search(config.pattern) !== -1) return config.averageCharsPerToken;
}
function getCharacterCount(text) {
return Array.from(text).length;
}
function estimateCjkTokens(segment) {
let kanaCount = 0;
let hangulCount = 0;
let hanziCount = 0;
for (const character of segment) {
const codePoint = character.codePointAt(0);
if (codePoint >= 12352 && codePoint <= 12543) kanaCount++;
else if (isHangulCodePoint(codePoint)) hangulCount++;
else hanziCount++;
}
return Math.ceil(hanziCount / HANZI_CHARS_PER_TOKEN + kanaCount / KANA_CHARS_PER_TOKEN + hangulCount / HANGUL_CHARS_PER_TOKEN);
}
function isHangulCodePoint(codePoint) {
return codePoint >= 44032 && codePoint <= 55215 || codePoint >= 4352 && codePoint <= 4607 || codePoint >= 12592 && codePoint <= 12687 || codePoint >= 43360 && codePoint <= 43391 || codePoint >= 55216 && codePoint <= 55295;
}
// src/index.ts
function isWithinTokenLimit(text, tokenLimit, options) {
let tokenCount = 0;
for (const segmentEstimate of walkSegments(text, options)) {
tokenCount += segmentEstimate.tokenCount;
if (tokenCount > tokenLimit) return false;
}
return tokenCount <= tokenLimit;
}
/** Estimates the number of tokens in a text string using heuristic rules. */
function estimateTokenCount(text, options = {}) {
if (!text) return 0;
let tokenCount = 0;
for (const segmentEstimate of walkSegments(text, options)) tokenCount += segmentEstimate.tokenCount;
return tokenCount;
}
/** Extracts a portion of text based on token positions, similar to `Array.prototype.slice()`. */
function sliceByTokens(text, start = 0, end, options = {}) {
if (!text) return "";
let segmentEstimates = walkSegments(text, options);
let totalTokens = 0;
if (start < 0 || end !== void 0 && end < 0) {
const bufferedEstimates = Array.from(segmentEstimates);
for (const { tokenCount } of bufferedEstimates) totalTokens += tokenCount;
segmentEstimates = bufferedEstimates;
}
const normalizedStart = start < 0 ? Math.max(0, totalTokens + start) : Math.max(0, start);
const normalizedEnd = end === void 0 ? 1e400 : end < 0 ? Math.max(0, totalTokens + end) : end;
if (normalizedStart >= normalizedEnd) return "";
const parts = [];
let currentTokenPos = 0;
for (const { segment, tokenCount } of segmentEstimates) {
if (currentTokenPos >= normalizedEnd) break;
const extracted = extractSegmentPart(segment, currentTokenPos, tokenCount, normalizedStart, normalizedEnd);
if (extracted) parts.push(extracted);
currentTokenPos += tokenCount;
}
return parts.join("");
}
function splitByTokens(text, tokensPerChunk, options = {}) {
if (!text || tokensPerChunk <= 0) return [];
const overlap = Math.min(options.overlap ?? 0, tokensPerChunk - 1);
const chunks = [];
let currentChunk = [];
let currentTokenCount = 0;
let hasUnchunkedSegments = false;
for (const segmentEstimate of walkSegments(text, options)) {
currentChunk.push(segmentEstimate);
currentTokenCount += segmentEstimate.tokenCount;
hasUnchunkedSegments = true;
if (currentTokenCount >= tokensPerChunk) {
chunks.push(joinSegments(currentChunk));
hasUnchunkedSegments = false;
if (overlap > 0) {
const overlapSegments = [];
let overlapTokenCount = 0;
for (let i = currentChunk.length - 1; i >= 0 && overlapTokenCount < overlap; i--) {
const overlapCandidate = currentChunk[i];
if (overlapTokenCount + overlapCandidate.tokenCount >= tokensPerChunk) break;
overlapSegments.unshift(overlapCandidate);
overlapTokenCount += overlapCandidate.tokenCount;
}
currentChunk = overlapSegments;
currentTokenCount = overlapTokenCount;
} else {
currentChunk = [];
currentTokenCount = 0;
}
}
}
if (currentChunk.length > 0 && hasUnchunkedSegments) chunks.push(joinSegments(currentChunk));
return chunks;
}
function joinSegments(segmentEstimates) {
let joined = "";
for (const { segment } of segmentEstimates) joined += segment;
return joined;
}
function extractSegmentPart(segment, segmentTokenStart, segmentTokenCount, targetStart, targetEnd) {
if (segmentTokenCount === 0) return segmentTokenStart >= targetStart && segmentTokenStart < targetEnd ? segment : "";
const segmentTokenEnd = segmentTokenStart + segmentTokenCount;
if (segmentTokenStart >= targetEnd || segmentTokenEnd <= targetStart) return "";
const overlapStart = Math.max(0, targetStart - segmentTokenStart);
const overlapEnd = Math.min(segmentTokenCount, targetEnd - segmentTokenStart);
if (overlapStart === 0 && overlapEnd === segmentTokenCount) return segment;
const characters = Array.from(segment);
const toCharacterIndex = (tokenOffset) => Math.round(tokenOffset / segmentTokenCount * characters.length);
return characters.slice(toCharacterIndex(overlapStart), toCharacterIndex(overlapEnd)).join("");
}
export { estimateTokenCount, isWithinTokenLimit, sliceByTokens, splitByTokens };