Remove split merging and add long document unit test

Author: dan-rubinstein

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

@@ -55,10 +55,7 @@ private List<ChunkOffset> chunk(String input, List<String> separators, int maxCh
             return chunkWithBackupChunker(input, maxChunkSize, chunkOffset);
         }
 
-        var potentialChunks = mergeChunkOffsetsUpToMaxChunkSize(
-            splitTextBySeparatorRegex(input, separators.get(separatorIndex)),
-            maxChunkSize
-        );
+        var potentialChunks = splitTextBySeparatorRegex(input, separators.get(separatorIndex));
         var actualChunks = new ArrayList<ChunkOffset>();
         for (var potentialChunk : potentialChunks) {
             if (isChunkWithinMaxSize(potentialChunk, maxChunkSize)) {
@@ -113,33 +110,6 @@ private List<ChunkOffsetAndCount> splitTextBySeparatorRegex(String input, String
         return chunkOffsets;
     }
 
-    private List<ChunkOffsetAndCount> mergeChunkOffsetsUpToMaxChunkSize(List<ChunkOffsetAndCount> chunkOffsets, int maxChunkSize) {
-        if (chunkOffsets.size() < 2) {
-            return chunkOffsets;
-        }
-
-        List<ChunkOffsetAndCount> mergedOffsetsAndCounts = new ArrayList<>();
-        var mergedChunk = chunkOffsets.getFirst();
-        for (int i = 1; i < chunkOffsets.size(); i++) {
-            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 {
-                mergedOffsetsAndCounts.add(mergedChunk);
-                mergedChunk = chunkOffsets.get(i);
-            }
-
-            if (i == chunkOffsets.size() - 1) {
-                mergedOffsetsAndCounts.add(mergedChunk);
-            }
-        }
-        return mergedOffsetsAndCounts;
-    }
-
     private List<ChunkOffset> chunkWithBackupChunker(String input, int maxChunkSize, int chunkOffset) {
         var chunks = new SentenceBoundaryChunker().chunk(input, new SentenceBoundaryChunkingSettings(maxChunkSize, 0));
         var chunksWithOffsets = new ArrayList<ChunkOffset>();

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

@@ -27,7 +27,17 @@ public static SeparatorSet fromString(String name) {
     public List<String> getSeparators() {
         return switch (this) {
             case PLAINTEXT -> List.of("\n\n", "\n");
-            case MARKDOWN -> List.of("\n#{1,6} ", "^(?!\\s*$).*\\n*{3,}\\n", "^(?!\\s*$).*\\n-{3,}\\n", "^(?!\\s*$).*\\n_{3,}\\n");
+            case MARKDOWN -> List.of(
+                "\n# ",
+                "\n## ",
+                "\n### ",
+                "\n#### ",
+                "\n##### ",
+                "\n###### ",
+                "^(?!\\s*$).*\\n*{3,}\\n",
+                "^(?!\\s*$).*\\n-{3,}\\n",
+                "^(?!\\s*$).*\\n_{3,}\\n"
+            );
         };
     }
 }

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.inference.ChunkingSettings;
 import org.elasticsearch.test.ESTestCase;
 
+import java.util.ArrayList;
 import java.util.List;
 
 public class RecursiveChunkerTests extends ESTestCase {
@@ -45,7 +46,7 @@ public void testChunkInputShorterThanMaxChunkSize() {
         assertExpectedChunksGenerated(input, settings, List.of(new Chunker.ChunkOffset(0, input.length())));
     }
 
-    public void testChunkInputRequiresOneSplitWithNoMerges() {
+    public void testChunkInputRequiresOneSplit() {
         List<String> separators = generateRandomSeparators();
         RecursiveChunkingSettings settings = generateChunkingSettings(10, separators);
         String input = generateTestText(2, List.of(separators.getFirst()));
@@ -57,23 +58,7 @@ public void testChunkInputRequiresOneSplitWithNoMerges() {
         );
     }
 
-    public void testChunkInputRequiresOneSplitWithMerges() {
-        List<String> separators = generateRandomSeparators();
-        RecursiveChunkingSettings settings = generateChunkingSettings(20, separators);
-        String input = generateTestText(3, List.of(separators.getFirst(), separators.getFirst()));
-
-        var expectedFirstChunkOffsetEnd = TEST_SENTENCE.length() * 2 + separators.getFirst().length();
-        assertExpectedChunksGenerated(
-            input,
-            settings,
-            List.of(
-                new Chunker.ChunkOffset(0, expectedFirstChunkOffsetEnd),
-                new Chunker.ChunkOffset(expectedFirstChunkOffsetEnd, input.length())
-            )
-        );
-    }
-
-    public void testChunkInputRequiresMultipleSplitsWithNoMerges() {
+    public void testChunkInputRequiresMultipleSplits() {
         var separators = generateRandomSeparators();
         RecursiveChunkingSettings settings = generateChunkingSettings(15, separators);
         String input = generateTestText(4, List.of(separators.get(1), separators.getFirst(), separators.get(1)));
@@ -93,22 +78,6 @@ public void testChunkInputRequiresMultipleSplitsWithNoMerges() {
         );
     }
 
-    public void testChunkInputRequiresMultipleSplitsWithMerges() {
-        var separators = generateRandomSeparators();
-        RecursiveChunkingSettings settings = generateChunkingSettings(25, separators);
-        String input = generateTestText(4, List.of(separators.get(1), separators.getFirst(), separators.get(1)));
-
-        var expectedFirstChunkOffsetEnd = TEST_SENTENCE.length() * 2 + separators.get(1).length();
-        assertExpectedChunksGenerated(
-            input,
-            settings,
-            List.of(
-                new Chunker.ChunkOffset(0, expectedFirstChunkOffsetEnd),
-                new Chunker.ChunkOffset(expectedFirstChunkOffsetEnd, input.length())
-            )
-        );
-    }
-
     public void testChunkInputRequiresBackupChunkingStrategy() {
         var separators = generateRandomSeparators();
         RecursiveChunkingSettings settings = generateChunkingSettings(10, separators);
@@ -149,6 +118,30 @@ public void testChunkWithRegexSeparator() {
         );
     }
 
+    public void testChunkLongDocument() {
+        int numSentences = randomIntBetween(50, 100);
+        List<String> separators = generateRandomSeparators();
+        List<String> splittersAfterSentences = new ArrayList<>();
+        for (int i = 0; i < numSentences - 1; i++) {
+            splittersAfterSentences.add(randomFrom(separators));
+        }
+        RecursiveChunkingSettings settings = generateChunkingSettings(15, separators);
+        String input = generateTestText(numSentences, splittersAfterSentences);
+
+        List<Chunker.ChunkOffset> expectedChunks = new ArrayList<>();
+        int currentOffset = 0;
+        for (int i = 0; i < numSentences; i++) {
+            int chunkLength = TEST_SENTENCE.length();
+            if (i > 0) {
+                chunkLength += splittersAfterSentences.get(i - 1).length();
+            }
+            expectedChunks.add(new Chunker.ChunkOffset(currentOffset, currentOffset + chunkLength));
+            currentOffset += chunkLength;
+        }
+
+        assertExpectedChunksGenerated(input, settings, expectedChunks);
+    }
+
     private void assertExpectedChunksGenerated(String input, RecursiveChunkingSettings settings, List<Chunker.ChunkOffset> expectedChunks) {
         RecursiveChunker chunker = new RecursiveChunker();
         List<Chunker.ChunkOffset> chunks = chunker.chunk(input, settings);