Add javadoc for chunker, add tests, reduce word counting operations

Author: dan-rubinstein

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/RecursiveChunker.java

@@ -16,6 +16,18 @@
 import java.util.List;
 import java.util.regex.Pattern;
 
+/**
+ * Split text into chunks recursively based on a list of separator regex strings.
+ * The maximum chunk size is measured in words and controlled
+ * by {@code maxNumberWordsPerChunk}. For each separator the chunker will go through the following process:
+ * 1. Split the text on each regex match of the separator.
+ * 2. Merge consecutive chunks when it is possible to do so without exceeding the max chunk size.
+ * 3. For each chunk after the merge:
+ *     1. Return it if it is within the maximum chunk size.
+ *     2. Repeat the process using the next separator in the list if the chunk exceeds the maximum chunk size.
+ *     If there are no more separators left to try, run the {@code SentenceBoundaryChunker} with the provided
+ *     max chunk size and no overlaps.
+ */
 public class RecursiveChunker implements Chunker {
     private final BreakIterator wordIterator;
 
@@ -35,7 +47,7 @@ public List<ChunkOffset> chunk(String input, ChunkingSettings chunkingSettings)
     }
 
     private List<ChunkOffset> chunk(String input, List<String> separators, int maxChunkSize, int separatorIndex, int chunkOffset) {
-        if (input.length() < 2 || isChunkWithinMaxSize(input, new ChunkOffset(0, input.length()), maxChunkSize)) {
+        if (input.length() < 2 || isChunkWithinMaxSize(buildChunkOffsetAndCount(input, 0, input.length()), maxChunkSize)) {
             return List.of(new ChunkOffset(chunkOffset, chunkOffset + input.length()));
         }
 
@@ -44,22 +56,23 @@ private List<ChunkOffset> chunk(String input, List<String> separators, int maxCh
         }
 
         var potentialChunks = mergeChunkOffsetsUpToMaxChunkSize(
-            input,
             splitTextBySeparatorRegex(input, separators.get(separatorIndex)),
             maxChunkSize
         );
         var actualChunks = new ArrayList<ChunkOffset>();
         for (var potentialChunk : potentialChunks) {
-            if (isChunkWithinMaxSize(input, potentialChunk, maxChunkSize)) {
-                actualChunks.add(new ChunkOffset(chunkOffset + potentialChunk.start(), chunkOffset + potentialChunk.end()));
+            if (isChunkWithinMaxSize(potentialChunk, maxChunkSize)) {
+                actualChunks.add(
+                    new ChunkOffset(chunkOffset + potentialChunk.chunkOffset.start(), chunkOffset + potentialChunk.chunkOffset.end())
+                );
             } else {
                 actualChunks.addAll(
                     chunk(
-                        input.substring(potentialChunk.start(), potentialChunk.end()),
+                        input.substring(potentialChunk.chunkOffset.start(), potentialChunk.chunkOffset.end()),
                         separators,
                         maxChunkSize,
                         separatorIndex + 1,
-                        chunkOffset + potentialChunk.start()
+                        chunkOffset + potentialChunk.chunkOffset.start()
                     )
                 );
             }
@@ -68,55 +81,63 @@ private List<ChunkOffset> chunk(String input, List<String> separators, int maxCh
         return actualChunks;
     }
 
-    private boolean isChunkWithinMaxSize(String fullText, ChunkOffset chunk, int maxChunkSize) {
+    private boolean isChunkWithinMaxSize(ChunkOffsetAndCount chunkOffsetAndCount, int maxChunkSize) {
+        return chunkOffsetAndCount.wordCount <= maxChunkSize;
+    }
+
+    private ChunkOffsetAndCount buildChunkOffsetAndCount(String fullText, int chunkStart, int chunkEnd) {
+        var chunkOffset = new ChunkOffset(chunkStart, chunkEnd);
+
         wordIterator.setText(fullText);
-        return ChunkerUtils.countWords(chunk.start(), chunk.end(), wordIterator) <= maxChunkSize;
+        return new ChunkOffsetAndCount(chunkOffset, ChunkerUtils.countWords(chunkStart, chunkEnd, wordIterator));
     }
 
-    private List<ChunkOffset> splitTextBySeparatorRegex(String input, String separatorRegex) {
+    private List<ChunkOffsetAndCount> splitTextBySeparatorRegex(String input, String separatorRegex) {
         var pattern = Pattern.compile(separatorRegex);
         var matcher = pattern.matcher(input);
 
-        var chunkOffsets = new ArrayList<ChunkOffset>();
+        var chunkOffsets = new ArrayList<ChunkOffsetAndCount>();
         int chunkStart = 0;
-        int searchStart = 0;
-        while (matcher.find(searchStart)) {
+        while (matcher.find()) {
             var chunkEnd = matcher.start();
             if (chunkStart < chunkEnd) {
-                chunkOffsets.add(new ChunkOffset(chunkStart, chunkEnd));
+                chunkOffsets.add(buildChunkOffsetAndCount(input, chunkStart, chunkEnd));
             }
             chunkStart = matcher.start();
-            searchStart = matcher.end();
         }
 
         if (chunkStart < input.length()) {
-            chunkOffsets.add(new ChunkOffset(chunkStart, input.length()));
+            chunkOffsets.add(buildChunkOffsetAndCount(input, chunkStart, input.length()));
         }
 
         return chunkOffsets;
     }
 
-    private List<ChunkOffset> mergeChunkOffsetsUpToMaxChunkSize(String input, List<ChunkOffset> chunkOffsets, int maxChunkSize) {
+    private List<ChunkOffsetAndCount> mergeChunkOffsetsUpToMaxChunkSize(List<ChunkOffsetAndCount> chunkOffsets, int maxChunkSize) {
         if (chunkOffsets.size() < 2) {
             return chunkOffsets;
         }
 
-        List<ChunkOffset> mergedOffsets = new ArrayList<>();
+        List<ChunkOffsetAndCount> mergedOffsetsAndCounts = new ArrayList<>();
         var mergedChunk = chunkOffsets.getFirst();
         for (int i = 1; i < chunkOffsets.size(); i++) {
-            var potentialMergedChunk = new ChunkOffset(mergedChunk.start(), chunkOffsets.get(i).end());
-            if (isChunkWithinMaxSize(input, potentialMergedChunk, maxChunkSize)) {
+            var chunkOffsetAndCountToMerge = chunkOffsets.get(i);
+            var potentialMergedChunk = new ChunkOffsetAndCount(
+                new ChunkOffset(mergedChunk.chunkOffset.start(), chunkOffsetAndCountToMerge.chunkOffset.end()),
+                mergedChunk.wordCount + chunkOffsetAndCountToMerge.wordCount
+            );
+            if (isChunkWithinMaxSize(potentialMergedChunk, maxChunkSize)) {
                 mergedChunk = potentialMergedChunk;
             } else {
-                mergedOffsets.add(mergedChunk);
+                mergedOffsetsAndCounts.add(mergedChunk);
                 mergedChunk = chunkOffsets.get(i);
             }
 
             if (i == chunkOffsets.size() - 1) {
-                mergedOffsets.add(mergedChunk);
+                mergedOffsetsAndCounts.add(mergedChunk);
             }
         }
-        return mergedOffsets;
+        return mergedOffsetsAndCounts;
     }
 
     private List<ChunkOffset> chunkWithBackupChunker(String input, int maxChunkSize, int chunkOffset) {
@@ -127,4 +148,6 @@ private List<ChunkOffset> chunkWithBackupChunker(String input, int maxChunkSize,
         }
         return chunksWithOffsets;
     }
+
+    private record ChunkOffsetAndCount(ChunkOffset chunkOffset, int wordCount) {}
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/RecursiveChunkerTests.java

@@ -129,6 +129,26 @@ public void testChunkInputRequiresBackupChunkingStrategy() {
         );
     }
 
+    public void testChunkWithRegexSeparator() {
+        var separators = List.of("(?<!\\n)\\n(?!\\n)", "(?<!\\n)\\n\\n(?!\\n)");
+        RecursiveChunkingSettings settings = generateChunkingSettings(10, separators);
+        String input = generateTestText(4, List.of("\n", "\n", "\n\n"));
+
+        var expectedFirstChunkOffsetEnd = TEST_SENTENCE.length();
+        var expectedSecondChunkOffsetEnd = TEST_SENTENCE.length() * 2 + "\n".length();
+        var expectedThirdChunkOffsetEnd = TEST_SENTENCE.length() * 3 + "\n".length() * 2;
+        assertExpectedChunksGenerated(
+            input,
+            settings,
+            List.of(
+                new Chunker.ChunkOffset(0, expectedFirstChunkOffsetEnd),
+                new Chunker.ChunkOffset(expectedFirstChunkOffsetEnd, expectedSecondChunkOffsetEnd),
+                new Chunker.ChunkOffset(expectedSecondChunkOffsetEnd, expectedThirdChunkOffsetEnd),
+                new Chunker.ChunkOffset(expectedThirdChunkOffsetEnd, input.length())
+            )
+        );
+    }
+
     private void assertExpectedChunksGenerated(String input, RecursiveChunkingSettings settings, List<Chunker.ChunkOffset> expectedChunks) {
         RecursiveChunker chunker = new RecursiveChunker();
         List<Chunker.ChunkOffset> chunks = chunker.chunk(input, settings);