chore: improve performance of sentence splitter (#2030)

.changeset/neat-sides-care.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+Improve performance of sentence splitter

examples/rag/longText.ts

@@ -1,3 +1,4 @@
+import { OpenAIEmbedding } from "@llamaindex/openai";
 import {
   Document,
   SentenceSplitter,
@@ -7,6 +8,10 @@ import {
 import { OldSentenceSplitter } from "./old-sentence-splitter";
 export const STORAGE_DIR = "./data";
 
+Settings.embedModel = new OpenAIEmbedding({
+  model: "text-embedding-3-small",
+});
+
 // Update node parser
 (async () => {
   // generate a document with a very long sentence (9000 words long)

packages/core/src/node-parser/sentence-splitter.ts

@@ -12,11 +12,7 @@ import {
   type TextSplitterFn,
 } from "./utils";
 
-type _Split = {
-  text: string;
-  isSentence: boolean;
-  tokenSize: number;
-};
+type _Split = [string, boolean, number]; // [text, isSentence, tokenSize]
 
 /**
  * Parse text with a preference for complete sentences.
@@ -114,34 +110,26 @@ export class SentenceSplitter extends MetadataAwareTextSplitter {
     return chunks;
   }
 
-  #split(text: string, chunkSize: number): _Split[] {
-    const tokenSize = this.tokenSize(text);
-    if (tokenSize <= chunkSize) {
-      return [
-        {
-          text,
-          isSentence: true,
-          tokenSize,
-        },
-      ];
+  *#split(
+    text: string,
+    chunkSize: number,
+    tokenSize?: number,
+  ): Generator<_Split> {
+    const _tokenSize = tokenSize ?? this.tokenSize(text);
+    if (_tokenSize <= chunkSize) {
+      yield [text, true, _tokenSize];
+      return;
     }
     const [textSplitsByFns, isSentence] = this.#getSplitsByFns(text);
-    const textSplits: _Split[] = [];
 
     for (const textSplit of textSplitsByFns) {
-      const tokenSize = this.tokenSize(textSplit);
-      if (tokenSize <= chunkSize) {
-        textSplits.push({
-          text: textSplit,
-          isSentence,
-          tokenSize,
-        });
+      const textSplitTokenSize = this.tokenSize(textSplit);
+      if (textSplitTokenSize <= chunkSize) {
+        yield [textSplit, isSentence, textSplitTokenSize];
       } else {
-        const recursiveTextSplits = this.#split(textSplit, chunkSize);
-        textSplits.push(...recursiveTextSplits);
+        yield* this.#split(textSplit, chunkSize, textSplitTokenSize);
       }
     }
-    return textSplits;
   }
 
   #getSplitsByFns(text: string): [splits: string[], isSentence: boolean] {
@@ -160,7 +148,7 @@ export class SentenceSplitter extends MetadataAwareTextSplitter {
     return [[text], true];
   }
 
-  #merge(splits: _Split[], chunkSize: number): string[] {
+  #merge(splits: Iterable<_Split>, chunkSize: number): string[] {
     const chunks: string[] = [];
     let currentChunk: [string, number][] = [];
     let lastChunk: [string, number][] = [];
@@ -186,23 +174,26 @@ export class SentenceSplitter extends MetadataAwareTextSplitter {
       }
     };
 
-    while (splits.length > 0) {
-      const curSplit = splits[0]!;
-      if (curSplit.tokenSize > chunkSize) {
+    const splitsIterator = splits[Symbol.iterator]();
+    let currentSplitResult = splitsIterator.next();
+
+    while (!currentSplitResult.done) {
+      const [text, isSentence, tokenSize] = currentSplitResult.value;
+      if (tokenSize > chunkSize) {
         throw new Error("Single token exceeded chunk size");
       }
-      if (currentChunkLength + curSplit.tokenSize > chunkSize && !newChunk) {
+      if (currentChunkLength + tokenSize > chunkSize && !newChunk) {
         closeChunk();
       } else {
         if (
-          curSplit.isSentence ||
-          currentChunkLength + curSplit.tokenSize <= chunkSize ||
+          isSentence ||
+          currentChunkLength + tokenSize <= chunkSize ||
           newChunk
         ) {
-          currentChunkLength += curSplit.tokenSize;
-          currentChunk.push([curSplit.text, curSplit.tokenSize]);
-          splits.shift();
+          currentChunkLength += tokenSize;
+          currentChunk.push([text, tokenSize]);
           newChunk = false;
+          currentSplitResult = splitsIterator.next();
         } else {
           closeChunk();
         }