embeddings-splitter
Version:
A typescript library to split your long texts into smaller chunks to send them to OpenAI Embeddings API
233 lines (232 loc) • 10.1 kB
JavaScript
"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
__setModuleDefault(result, mod);
return result;
};
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.TokenTextSplitter = exports.RecursiveCharacterTextSplitter = exports.CharacterTextSplitter = exports.TextSplitter = void 0;
const document_js_1 = require("./document.js");
class TextSplitter {
constructor(fields) {
var _a, _b;
this.chunkSize = 1000;
this.chunkOverlap = 200;
this.chunkSize = (_a = fields === null || fields === void 0 ? void 0 : fields.chunkSize) !== null && _a !== void 0 ? _a : this.chunkSize;
this.chunkOverlap = (_b = fields === null || fields === void 0 ? void 0 : fields.chunkOverlap) !== null && _b !== void 0 ? _b : this.chunkOverlap;
if (this.chunkOverlap >= this.chunkSize) {
throw new Error("Cannot have chunkOverlap >= chunkSize");
}
}
createDocuments(texts,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
metadatas = []) {
return __awaiter(this, void 0, void 0, function* () {
const _metadatas = metadatas.length > 0 ? metadatas : new Array(texts.length).fill({});
const documents = new Array();
for (let i = 0; i < texts.length; i += 1) {
const text = texts[i];
for (const chunk of yield this.splitText(text)) {
documents.push(new document_js_1.Document({ pageContent: chunk, metadata: _metadatas[i] }));
}
}
return documents;
});
}
splitDocuments(documents) {
return __awaiter(this, void 0, void 0, function* () {
const texts = documents.map((doc) => doc.pageContent);
const metadatas = documents.map((doc) => doc.metadata);
return this.createDocuments(texts, metadatas);
});
}
joinDocs(docs, separator) {
const text = docs.join(separator).trim();
return text === "" ? null : text;
}
mergeSplits(splits, separator) {
const docs = [];
const currentDoc = [];
let total = 0;
for (const d of splits) {
const _len = d.length;
if (total + _len >= this.chunkSize) {
if (total > this.chunkSize) {
console.warn(`Created a chunk of size ${total}, +
which is longer than the specified ${this.chunkSize}`);
}
if (currentDoc.length > 0) {
const doc = this.joinDocs(currentDoc, separator);
if (doc !== null) {
docs.push(doc);
}
// Keep on popping if:
// - we have a larger chunk than in the chunk overlap
// - or if we still have any chunks and the length is long
while (total > this.chunkOverlap ||
(total + _len > this.chunkSize && total > 0)) {
total -= currentDoc[0].length;
currentDoc.shift();
}
}
}
currentDoc.push(d);
total += _len;
}
const doc = this.joinDocs(currentDoc, separator);
if (doc !== null) {
docs.push(doc);
}
return docs;
}
}
exports.TextSplitter = TextSplitter;
class CharacterTextSplitter extends TextSplitter {
constructor(fields) {
var _a;
super(fields);
this.separator = "\n\n";
this.separator = (_a = fields === null || fields === void 0 ? void 0 : fields.separator) !== null && _a !== void 0 ? _a : this.separator;
}
splitText(text) {
return __awaiter(this, void 0, void 0, function* () {
// First we naively split the large input into a bunch of smaller ones.
let splits;
if (this.separator) {
splits = text.split(this.separator);
}
else {
splits = text.split("");
}
return this.mergeSplits(splits, this.separator);
});
}
}
exports.CharacterTextSplitter = CharacterTextSplitter;
class RecursiveCharacterTextSplitter extends TextSplitter {
constructor(fields) {
var _a;
super(fields);
this.separators = ["\n\n", "\n", " ", ""];
this.separators = (_a = fields === null || fields === void 0 ? void 0 : fields.separators) !== null && _a !== void 0 ? _a : this.separators;
}
splitText(text) {
return __awaiter(this, void 0, void 0, function* () {
const finalChunks = [];
// Get appropriate separator to use
let separator = this.separators[this.separators.length - 1];
for (const s of this.separators) {
if (s === "") {
separator = s;
break;
}
if (text.includes(s)) {
separator = s;
break;
}
}
// Now that we have the separator, split the text
let splits;
if (separator) {
splits = text.split(separator);
}
else {
splits = text.split("");
}
// Now go merging things, recursively splitting longer texts.
let goodSplits = [];
for (const s of splits) {
if (s.length < this.chunkSize) {
goodSplits.push(s);
}
else {
if (goodSplits.length) {
const mergedText = this.mergeSplits(goodSplits, separator);
finalChunks.push(...mergedText);
goodSplits = [];
}
const otherInfo = yield this.splitText(s);
finalChunks.push(...otherInfo);
}
}
if (goodSplits.length) {
const mergedText = this.mergeSplits(goodSplits, separator);
finalChunks.push(...mergedText);
}
return finalChunks;
});
}
}
exports.RecursiveCharacterTextSplitter = RecursiveCharacterTextSplitter;
/**
* Implementation of splitter which looks at tokens.
*/
class TokenTextSplitter extends TextSplitter {
constructor(fields) {
var _a, _b, _c;
super(fields);
this.encodingName = (_a = fields === null || fields === void 0 ? void 0 : fields.encodingName) !== null && _a !== void 0 ? _a : "gpt2";
this.allowedSpecial = (_b = fields === null || fields === void 0 ? void 0 : fields.allowedSpecial) !== null && _b !== void 0 ? _b : [];
this.disallowedSpecial = (_c = fields === null || fields === void 0 ? void 0 : fields.disallowedSpecial) !== null && _c !== void 0 ? _c : "all";
}
splitText(text) {
return __awaiter(this, void 0, void 0, function* () {
if (!this.tokenizer) {
const tiktoken = yield TokenTextSplitter.imports();
this.tokenizer = tiktoken.get_encoding(this.encodingName);
}
const splits = [];
const input_ids = this.tokenizer.encode(text, this.allowedSpecial, this.disallowedSpecial);
let start_idx = 0;
let cur_idx = Math.min(start_idx + this.chunkSize, input_ids.length);
let chunk_ids = input_ids.slice(start_idx, cur_idx);
const decoder = new TextDecoder();
while (start_idx < input_ids.length) {
splits.push(decoder.decode(this.tokenizer.decode(chunk_ids)));
start_idx += this.chunkSize - this.chunkOverlap;
cur_idx = Math.min(start_idx + this.chunkSize, input_ids.length);
chunk_ids = input_ids.slice(start_idx, cur_idx);
}
return splits;
});
}
static imports() {
return __awaiter(this, void 0, void 0, function* () {
try {
return yield Promise.resolve().then(() => __importStar(require("@dqbd/tiktoken")));
}
catch (err) {
console.error(err);
throw new Error("Please install @dqbd/tiktoken as a dependency with, e.g. `npm install -S @dqbd/tiktoken`");
}
});
}
}
exports.TokenTextSplitter = TokenTextSplitter;