chonkie
Version:
🦛 CHONK your texts in TS with Chonkie!✨The no-nonsense lightweight and efficient chunking library.
458 lines • 21.6 kB
JavaScript
;
/** Module containing RecursiveChunker class. */
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.RecursiveChunker = void 0;
const tokenizer_1 = require("../tokenizer");
const recursive_1 = require("../types/recursive");
const base_1 = require("./base");
const hub_1 = require("../utils/hub");
/**
* Recursively chunk text using a set of rules.
*
* This class extends the BaseChunker class and implements the chunk method.
* It provides a flexible way to chunk text based on custom rules, including
* delimiters, whitespace, and token-based chunking.
*
* @extends BaseChunker
* @property {number} chunkSize - The maximum number of tokens per chunk.
* @property {number} minCharactersPerChunk - The minimum number of characters per chunk.
* @property {RecursiveRules} rules - The rules that define how text should be recursively chunked.
* @property {string} sep - The separator string used for internal splitting (usually "✄").
*
* @method chunk - Recursively chunk a single text into chunks or strings.
* @method chunkBatch - Recursively chunk a batch of texts.
* @method toString - Returns a string representation of the RecursiveChunker instance.
* @method call - Call the chunker with a single string or an array of strings. (see callable signatures)
*
* @static
* @method create
* @memberof RecursiveChunker
* @param {RecursiveChunkerOptions} [options] - Configuration options for the RecursiveChunker.
* @returns {Promise<RecursiveChunker>} A Promise that resolves to a RecursiveChunker instance.
*
* @example
* const chunker = await RecursiveChunker.create({ chunkSize: 256 });
* const chunks = await chunker("Some text to chunk");
* const batchChunks = await chunker(["Text 1", "Text 2"]);
*/
class RecursiveChunker extends base_1.BaseChunker {
/**
* Private constructor. Use `RecursiveChunker.create()` to instantiate.
*/
constructor(tokenizer, chunkSize, rules, minCharactersPerChunk) {
super(tokenizer);
this._CHARS_PER_TOKEN = 6.5;
if (chunkSize <= 0) {
throw new Error("chunkSize must be greater than 0");
}
if (minCharactersPerChunk <= 0) {
throw new Error("minCharactersPerChunk must be greater than 0");
}
if (!(rules instanceof recursive_1.RecursiveRules)) {
throw new Error("rules must be a RecursiveRules object");
}
this.chunkSize = chunkSize;
this.minCharactersPerChunk = minCharactersPerChunk;
this.rules = rules;
this.sep = "✄";
}
/**
* Creates and initializes a directly callable RecursiveChunker instance.
*
* This static factory method constructs a RecursiveChunker with the provided options and returns a callable function object.
* The returned instance can be used as both a function (to chunk text(s)) and as an object (with all RecursiveChunker methods and properties).
*
* @param {RecursiveChunkerOptions} [options] - Configuration options for the chunker. All options are optional:
* @param {string|Tokenizer} [options.tokenizer="Xenova/gpt2"] - Tokenizer to use for text processing. Can be a string identifier (e.g., "Xenova/gpt2") or a Tokenizer instance. If a string is provided, Tokenizer.create() is called internally.
* @param {number} [options.chunkSize=512] - Maximum number of tokens per chunk. Must be > 0.
* @param {RecursiveRules} [options.rules=new RecursiveRules()] - Rules for recursive chunking. See {@link RecursiveRules} for customization.
* @param {number} [options.minCharactersPerChunk=24] - Minimum number of characters per chunk. Must be > 0.
*
* @returns {Promise<CallableRecursiveChunker>} Promise resolving to a callable RecursiveChunker instance.
*
* @throws {Error} If any option is invalid (e.g., chunkSize <= 0).
*
* @see CallableRecursiveChunker for the callable interface and available properties/methods.
*
* @example <caption>Basic usage with default options</caption>
* const chunker = await RecursiveChunker.create();
* const chunks = await chunker("Some text to chunk");
*
* @example <caption>Custom options and batch chunking</caption>
* const chunker = await RecursiveChunker.create({ chunkSize: 256 });
* const batchChunks = await chunker(["Text 1", "Text 2"]);
*
* @example <caption>Accessing properties and methods</caption>
* const chunker = await RecursiveChunker.create();
* console.log(chunker.chunkSize); // 512
* console.log(chunker.rules); // RecursiveRules instance
* const chunks = await chunker.chunk("Some text"); // Use as object method
*
* @note
* The returned instance is both callable (like a function) and has all properties/methods of RecursiveChunker.
* You can use it as a drop-in replacement for a function or a class instance.
*
* @note
* For advanced customization, pass a custom RecursiveRules object to the rules option.
* See {@link RecursiveRules} and {@link RecursiveLevel} for rule structure.
*/
static create() {
return __awaiter(this, arguments, void 0, function* (options = {}) {
const { tokenizer = "Xenova/gpt2", chunkSize = 512, rules = new recursive_1.RecursiveRules(), minCharactersPerChunk = 24, } = options;
let tokenizerInstance;
if (typeof tokenizer === 'string') {
tokenizerInstance = yield tokenizer_1.Tokenizer.create(tokenizer);
}
else {
tokenizerInstance = tokenizer;
}
const plainInstance = new RecursiveChunker(tokenizerInstance, chunkSize, rules, minCharactersPerChunk);
// Create the callable function wrapper
const callableFn = function (textOrTexts, showProgress) {
if (typeof textOrTexts === 'string') {
return plainInstance.call(textOrTexts, showProgress);
}
else {
return plainInstance.call(textOrTexts, showProgress);
}
};
// Set the prototype so that 'instanceof RecursiveChunker' works
Object.setPrototypeOf(callableFn, RecursiveChunker.prototype);
// Copy all enumerable own properties from plainInstance to callableFn
Object.assign(callableFn, plainInstance);
return callableFn;
});
}
/**
* Creates and initializes a RecursiveChunker instance from a recipe that is directly callable.
*
* This method loads a recipe from the Chonkie hub and uses the recipe's recursive rules
* to configure the RecursiveChunker. The recipe rules override the default rules.
*
* @param {RecursiveChunkerRecipeOptions} [options] - Options for configuring the RecursiveChunker with recipe settings.
* @returns {Promise<CallableRecursiveChunker>} A promise that resolves to a callable RecursiveChunker instance.
*
* @example
* const chunker = await RecursiveChunker.fromRecipe({ name: 'default', language: 'en' });
* const chunks = await chunker("This is a sample text.");
*
* @see RecursiveChunkerRecipeOptions
*/
static fromRecipe() {
return __awaiter(this, arguments, void 0, function* (options = {}) {
var _a, _b;
const { name = 'default', language = 'en', filePath, tokenizer = "Xenova/gpt2", chunkSize = 512, minCharactersPerChunk = 24 } = options;
// Load the recipe using Hubbie
const hubbie = new hub_1.Hubbie();
const recipe = yield hubbie.getRecipe(name, language, filePath);
// Extract recursive_rules from recipe and create RecursiveRules
let rules;
if ((_b = (_a = recipe.recipe) === null || _a === void 0 ? void 0 : _a.recursive_rules) === null || _b === void 0 ? void 0 : _b.levels) {
const levels = recipe.recipe.recursive_rules.levels.map((levelData) => {
// Map from recipe format to RecursiveLevel format
const levelConfig = {
whitespace: levelData.whitespace || false,
includeDelim: levelData.include_delim || 'prev'
};
// Only add delimiters if they exist and are not null
if (levelData.delimiters != null) {
levelConfig.delimiters = levelData.delimiters;
}
return new recursive_1.RecursiveLevel(levelConfig);
});
rules = new recursive_1.RecursiveRules({ levels: levels.map((level) => level.toDict()) });
}
else {
// Fallback to default rules if no recursive_rules in recipe
rules = new recursive_1.RecursiveRules();
}
// Create the RecursiveChunker using the regular create method with recipe values
return RecursiveChunker.create({
tokenizer,
chunkSize,
rules,
minCharactersPerChunk
});
});
}
/**
* Estimates the number of tokens in a given text.
*
* This method uses a character-to-token ratio (default: 6.5 characters per token) for quick estimation.
* If the estimated token count exceeds the chunk size, it performs an actual token count.
*
* @param {string} text - The text to estimate token count for
* @returns {Promise<number>} A promise that resolves to the estimated number of tokens
* @private
*/
_estimateTokenCount(text) {
return __awaiter(this, void 0, void 0, function* () {
const estimate = Math.max(1, Math.floor(text.length / this._CHARS_PER_TOKEN));
return estimate > this.chunkSize ? this.chunkSize + 1 : yield this.tokenizer.countTokens(text);
});
}
/**
* Split the text into chunks based on the provided recursive level rules.
*
* This method handles three different splitting strategies:
* 1. Whitespace-based splitting: Splits text on spaces
* 2. Delimiter-based splitting: Splits text on specified delimiters with options to include delimiters
* 3. Token-based splitting: Splits text into chunks of maximum token size
*
* @param {string} text - The text to be split into chunks
* @param {RecursiveLevel} recursiveLevel - The rules defining how to split the text
* @returns {Promise<string[]>} A promise that resolves to an array of text chunks
* @private
*/
_splitText(text, recursiveLevel) {
return __awaiter(this, void 0, void 0, function* () {
// At every delimiter, replace it with the sep
if (recursiveLevel.whitespace) {
return text.split(" ");
}
else if (recursiveLevel.delimiters) {
let t = text;
if (recursiveLevel.includeDelim === "prev") {
for (const delimiter of Array.isArray(recursiveLevel.delimiters) ? recursiveLevel.delimiters : [recursiveLevel.delimiters]) {
t = t.replace(delimiter, delimiter + this.sep);
}
}
else if (recursiveLevel.includeDelim === "next") {
for (const delimiter of Array.isArray(recursiveLevel.delimiters) ? recursiveLevel.delimiters : [recursiveLevel.delimiters]) {
t = t.replace(delimiter, this.sep + delimiter);
}
}
else {
for (const delimiter of Array.isArray(recursiveLevel.delimiters) ? recursiveLevel.delimiters : [recursiveLevel.delimiters]) {
t = t.replace(delimiter, this.sep);
}
}
const splits = t.split(this.sep).filter(split => split !== "");
// Merge short splits
let current = "";
const merged = [];
for (const split of splits) {
if (split.length < this.minCharactersPerChunk) {
current += split;
}
else if (current) {
current += split;
merged.push(current);
current = "";
}
else {
merged.push(split);
}
if (current.length >= this.minCharactersPerChunk) {
merged.push(current);
current = "";
}
}
if (current) {
merged.push(current);
}
return merged;
}
else {
// Encode, Split, and Decode
const encoded = yield this.tokenizer.encode(text);
const tokenSplits = [];
for (let i = 0; i < encoded.length; i += this.chunkSize) {
tokenSplits.push(encoded.slice(i, i + this.chunkSize));
}
return yield this.tokenizer.decodeBatch(tokenSplits);
}
});
}
/**
* Create a RecursiveChunk object with indices based on the current offset.
*
* This method constructs a RecursiveChunk object that contains metadata about the chunk,
* including the text content, its start and end indices, token count, and the level of recursion.
*
* @param {string} text - The text content of the chunk
* @param {number} tokenCount - The number of tokens in the chunk
*/
_makeChunks(text, tokenCount, level, startOffset) {
return new recursive_1.RecursiveChunk({
text: text,
startIndex: startOffset,
endIndex: startOffset + text.length,
tokenCount: tokenCount,
level: level
});
}
/**
* Merge short splits.
*/
_mergeSplits(splits, tokenCounts, combineWhitespace = false) {
if (!splits.length || !tokenCounts.length) {
return [[], []];
}
// If the number of splits and token counts does not match, raise an error
if (splits.length !== tokenCounts.length) {
throw new Error(`Number of splits ${splits.length} does not match number of token counts ${tokenCounts.length}`);
}
// If all splits are larger than the chunk size, return them
if (tokenCounts.every(count => count > this.chunkSize)) {
return [splits, tokenCounts];
}
// If the splits are too short, merge them
const merged = [];
const cumulativeTokenCounts = [];
let sum = 0;
if (combineWhitespace) {
// +1 for the whitespace
cumulativeTokenCounts.push(0);
for (const count of tokenCounts) {
sum += count + 1;
cumulativeTokenCounts.push(sum);
}
}
else {
cumulativeTokenCounts.push(0);
for (const count of tokenCounts) {
sum += count;
cumulativeTokenCounts.push(sum);
}
}
let currentIndex = 0;
const combinedTokenCounts = [];
while (currentIndex < splits.length) {
const currentTokenCount = cumulativeTokenCounts[currentIndex];
const requiredTokenCount = currentTokenCount + this.chunkSize;
// Find the index to merge at
let index = this._bisectLeft(cumulativeTokenCounts, requiredTokenCount, currentIndex) - 1;
index = Math.min(index, splits.length);
// If currentIndex == index, we need to move to the next index
if (index === currentIndex) {
index += 1;
}
// Merge splits
if (combineWhitespace) {
merged.push(splits.slice(currentIndex, index).join(" "));
}
else {
merged.push(splits.slice(currentIndex, index).join(""));
}
// Adjust token count
combinedTokenCounts.push(cumulativeTokenCounts[Math.min(index, splits.length)] - currentTokenCount);
currentIndex = index;
}
return [merged, combinedTokenCounts];
}
/**
* Binary search to find the leftmost position where value should be inserted to maintain order.
*
* @param {number[]} arr - The array to search
* @param {number} value - The value to insert
* @param {number} [lo=0] - The starting index for the search
* @returns {number} The index where the value should be inserted
* @private
*/
_bisectLeft(arr, value, lo = 0) {
let hi = arr.length;
while (lo < hi) {
const mid = (lo + hi) >>> 1;
if (arr[mid] < value) {
lo = mid + 1;
}
else {
hi = mid;
}
}
return lo;
}
/**
* Recursive helper for core chunking.
*/
_recursiveChunk(text_1) {
return __awaiter(this, arguments, void 0, function* (text, level = 0, startOffset = 0) {
if (!text) {
return [];
}
if (level >= this.rules.length) {
const tokenCount = yield this._estimateTokenCount(text);
return [
this._makeChunks(text, tokenCount, level, startOffset)
];
}
const currRule = this.rules.getLevel(level);
if (!currRule) {
throw new Error(`No rule found at level ${level}`);
}
const splits = yield this._splitText(text, currRule);
const tokenCounts = yield Promise.all(splits.map(split => this._estimateTokenCount(split)));
let merged;
let combinedTokenCounts;
if (currRule.delimiters === undefined && !currRule.whitespace) {
[merged, combinedTokenCounts] = [splits, tokenCounts];
}
else if (currRule.delimiters === undefined && currRule.whitespace) {
[merged, combinedTokenCounts] = this._mergeSplits(splits, tokenCounts, true);
// NOTE: This is a hack to fix the reconstruction issue when whitespace is used.
// When whitespace is there, " ".join only adds space between words, not before the first word.
// To make it combine back properly, all splits except the first one are prefixed with a space.
merged = merged.slice(0, 1).concat(merged.slice(1).map(text => " " + text));
}
else {
[merged, combinedTokenCounts] = this._mergeSplits(splits, tokenCounts, false);
}
// Chunk long merged splits
const chunks = [];
let currentOffset = startOffset;
for (let i = 0; i < merged.length; i++) {
const split = merged[i];
const tokenCount = combinedTokenCounts[i];
if (tokenCount > this.chunkSize) {
chunks.push(...yield this._recursiveChunk(split, level + 1, currentOffset));
}
else {
chunks.push(this._makeChunks(split, tokenCount, level, currentOffset));
}
// Update the offset by the length of the processed split.
currentOffset += split.length;
}
return chunks;
});
}
/**
* Recursively chunk text.
*
* This method is the main entry point for chunking text using the RecursiveChunker.
* It takes a single text string and returns an array of RecursiveChunk objects.
*
* @param {string} text - The text to be chunked
* @returns {Promise<RecursiveChunk[]>} A promise that resolves to an array of RecursiveChunk objects
*/
chunk(text) {
return __awaiter(this, void 0, void 0, function* () {
const result = yield this._recursiveChunk(text, 0, 0);
return result;
});
}
/**
* Return a string representation of the RecursiveChunker.
*
* This method provides a string representation of the RecursiveChunker instance,
* including its tokenizer, rules, chunk size, minimum characters per chunk, and return type.
*
* @returns {string} A string representation of the RecursiveChunker
*/
toString() {
return `RecursiveChunker(tokenizer=${this.tokenizer}, ` +
`rules=${this.rules}, chunkSize=${this.chunkSize}, ` +
`minCharactersPerChunk=${this.minCharactersPerChunk})`;
}
}
exports.RecursiveChunker = RecursiveChunker;
//# sourceMappingURL=recursive.js.map