Add recursive chunker

docs/changelog/126866.yaml

@@ -0,0 +1,5 @@
+pr: 126866
+summary: Add recursive chunker
+area: Machine Learning
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/inference/ChunkingStrategy.java

@@ -15,7 +15,8 @@
 
 public enum ChunkingStrategy {
     WORD("word"),
-    SENTENCE("sentence");
+    SENTENCE("sentence"),
+    RECURSIVE("recursive");
 
     private final String chunkingStrategy;
 

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

@@ -26,6 +26,7 @@
 import org.elasticsearch.xpack.core.inference.results.TextEmbeddingByteResults;
 import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
 import org.elasticsearch.xpack.inference.action.task.StreamingTaskManager;
+import org.elasticsearch.xpack.inference.chunking.RecursiveChunkingSettings;
 import org.elasticsearch.xpack.inference.chunking.SentenceBoundaryChunkingSettings;
 import org.elasticsearch.xpack.inference.chunking.WordBoundaryChunkingSettings;
 import org.elasticsearch.xpack.inference.common.amazon.AwsSecretSettings;
@@ -472,6 +473,9 @@ private static void addChunkingSettingsNamedWriteables(List<NamedWriteableRegist
                 SentenceBoundaryChunkingSettings::new
             )
         );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(ChunkingSettings.class, RecursiveChunkingSettings.NAME, RecursiveChunkingSettings::new)
+        );
     }
 
     private static void addInferenceResultsNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {

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

@@ -18,6 +18,7 @@ public static Chunker fromChunkingStrategy(ChunkingStrategy chunkingStrategy) {
         return switch (chunkingStrategy) {
             case WORD -> new WordBoundaryChunker();
             case SENTENCE -> new SentenceBoundaryChunker();
+            case RECURSIVE -> new RecursiveChunker();
         };
     }
 }

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

@@ -0,0 +1,31 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.chunking;
+
+import com.ibm.icu.text.BreakIterator;
+
+public class ChunkerUtils {
+
+    // setText() should be applied before using this function.
+    static int countWords(int start, int end, BreakIterator wordIterator) {
+        assert start < end;
+        wordIterator.preceding(start); // start of the current word
+
+        int boundary = wordIterator.current();
+        int wordCount = 0;
+        while (boundary != BreakIterator.DONE && boundary <= end) {
+            int wordStatus = wordIterator.getRuleStatus();
+            if (wordStatus != BreakIterator.WORD_NONE) {
+                wordCount++;
+            }
+            boundary = wordIterator.next();
+        }
+
+        return wordCount;
+    }
+}

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

@@ -47,6 +47,7 @@ public static ChunkingSettings fromMap(Map<String, Object> settings, boolean ret
         return switch (chunkingStrategy) {
             case WORD -> WordBoundaryChunkingSettings.fromMap(new HashMap<>(settings));
             case SENTENCE -> SentenceBoundaryChunkingSettings.fromMap(new HashMap<>(settings));
+            case RECURSIVE -> RecursiveChunkingSettings.fromMap(new HashMap<>(settings));
         };
     }
 }

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

@@ -11,7 +11,9 @@ public enum ChunkingSettingsOptions {
     STRATEGY("strategy"),
     MAX_CHUNK_SIZE("max_chunk_size"),
     OVERLAP("overlap"),
-    SENTENCE_OVERLAP("sentence_overlap");
+    SENTENCE_OVERLAP("sentence_overlap"),
+    SEPARATOR_SET("separator_set"),
+    SEPARATORS("separators");
 
     private final String chunkingSettingsOption;
 

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

@@ -0,0 +1,153 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.chunking;
+
+import com.ibm.icu.text.BreakIterator;
+
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.inference.ChunkingSettings;
+
+import java.util.ArrayList;
+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;
+
+    public RecursiveChunker() {
+        wordIterator = BreakIterator.getWordInstance();
+    }
+
+    @Override
+    public List<ChunkOffset> chunk(String input, ChunkingSettings chunkingSettings) {
+        if (chunkingSettings instanceof RecursiveChunkingSettings recursiveChunkingSettings) {
+            return chunk(input, recursiveChunkingSettings.getSeparators(), recursiveChunkingSettings.getMaxChunkSize(), 0, 0);
+        } else {
+            throw new IllegalArgumentException(
+                Strings.format("RecursiveChunker can't use ChunkingSettings with strategy [%s]", chunkingSettings.getChunkingStrategy())
+            );
+        }
+    }
+
+    private List<ChunkOffset> chunk(String input, List<String> separators, int maxChunkSize, int separatorIndex, int chunkOffset) {
+        if (input.length() < 2 || isChunkWithinMaxSize(buildChunkOffsetAndCount(input, 0, input.length()), maxChunkSize)) {
+            return List.of(new ChunkOffset(chunkOffset, chunkOffset + input.length()));
+        }
+
+        if (separatorIndex > separators.size() - 1) {
+            return chunkWithBackupChunker(input, maxChunkSize, chunkOffset);
+        }
+
+        var potentialChunks = mergeChunkOffsetsUpToMaxChunkSize(
+            splitTextBySeparatorRegex(input, separators.get(separatorIndex)),
+            maxChunkSize
+        );
+        var actualChunks = new ArrayList<ChunkOffset>();
+        for (var potentialChunk : potentialChunks) {
+            if (isChunkWithinMaxSize(potentialChunk, maxChunkSize)) {
+                actualChunks.add(
+                    new ChunkOffset(chunkOffset + potentialChunk.chunkOffset.start(), chunkOffset + potentialChunk.chunkOffset.end())
+                );
+            } else {
+                actualChunks.addAll(
+                    chunk(
+                        input.substring(potentialChunk.chunkOffset.start(), potentialChunk.chunkOffset.end()),
+                        separators,
+                        maxChunkSize,
+                        separatorIndex + 1,
+                        chunkOffset + potentialChunk.chunkOffset.start()
+                    )
+                );
+            }
+        }
+
+        return actualChunks;
+    }
+
+    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 new ChunkOffsetAndCount(chunkOffset, ChunkerUtils.countWords(chunkStart, chunkEnd, wordIterator));
+    }
+
+    private List<ChunkOffsetAndCount> splitTextBySeparatorRegex(String input, String separatorRegex) {
+        var pattern = Pattern.compile(separatorRegex);
+        var matcher = pattern.matcher(input);
+
+        var chunkOffsets = new ArrayList<ChunkOffsetAndCount>();
+        int chunkStart = 0;
+        while (matcher.find()) {
+            var chunkEnd = matcher.start();
+            if (chunkStart < chunkEnd) {
+                chunkOffsets.add(buildChunkOffsetAndCount(input, chunkStart, chunkEnd));
+            }
+            chunkStart = matcher.start();
+        }
+
+        if (chunkStart < input.length()) {
+            chunkOffsets.add(buildChunkOffsetAndCount(input, chunkStart, input.length()));
+        }
+
+        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>();
+        for (var chunk : chunks) {
+            chunksWithOffsets.add(new ChunkOffset(chunk.start() + chunkOffset, chunk.end() + chunkOffset));
+        }
+        return chunksWithOffsets;
+    }
+
+    private record ChunkOffsetAndCount(ChunkOffset chunkOffset, int wordCount) {}
+}