[ML] Add sentence overlap option to the sentence chunking settings (elastic#114461) (elastic#114626)

Author: davidkyle

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -240,7 +240,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion SIMULATE_INDEX_TEMPLATES_SUBSTITUTIONS = def(8_764_00_0);
     public static final TransportVersion RETRIEVERS_TELEMETRY_ADDED = def(8_765_00_0);
     public static final TransportVersion ESQL_CACHED_STRING_SERIALIZATION = def(8_766_00_0);
-    public static final TransportVersion CHUNK_SENTENCE_OVERLAP_SETTING_ADDE1D = def(8_767_00_0);
+    public static final TransportVersion CHUNK_SENTENCE_OVERLAP_SETTING_ADDED = def(8_767_00_0);
     public static final TransportVersion OPT_IN_ESQL_CCS_EXECUTION_INFO = def(8_768_00_0);
 
     /*

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

@@ -10,7 +10,8 @@
 public enum ChunkingSettingsOptions {
     STRATEGY("strategy"),
     MAX_CHUNK_SIZE("max_chunk_size"),
-    OVERLAP("overlap");
+    OVERLAP("overlap"),
+    SENTENCE_OVERLAP("sentence_overlap");
 
     private final String chunkingSettingsOption;
 

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

@@ -34,6 +34,7 @@ public class SentenceBoundaryChunker implements Chunker {
     public SentenceBoundaryChunker() {
         sentenceIterator = BreakIterator.getSentenceInstance(Locale.ROOT);
         wordIterator = BreakIterator.getWordInstance(Locale.ROOT);
+
     }
 
     /**
@@ -46,7 +47,7 @@ public SentenceBoundaryChunker() {
     @Override
     public List<String> chunk(String input, ChunkingSettings chunkingSettings) {
         if (chunkingSettings instanceof SentenceBoundaryChunkingSettings sentenceBoundaryChunkingSettings) {
-            return chunk(input, sentenceBoundaryChunkingSettings.maxChunkSize);
+            return chunk(input, sentenceBoundaryChunkingSettings.maxChunkSize, sentenceBoundaryChunkingSettings.sentenceOverlap > 0);
         } else {
             throw new IllegalArgumentException(
                 Strings.format(
@@ -64,7 +65,7 @@ public List<String> chunk(String input, ChunkingSettings chunkingSettings) {
      * @param maxNumberWordsPerChunk Maximum size of the chunk
      * @return The input text chunked
      */
-    public List<String> chunk(String input, int maxNumberWordsPerChunk) {
+    public List<String> chunk(String input, int maxNumberWordsPerChunk, boolean includePrecedingSentence) {
         var chunks = new ArrayList<String>();
 
         sentenceIterator.setText(input);
@@ -75,24 +76,46 @@ public List<String> chunk(String input, int maxNumberWordsPerChunk) {
         int sentenceStart = 0;
         int chunkWordCount = 0;
 
+        int wordsInPrecedingSentenceCount = 0;
+        int previousSentenceStart = 0;
+
         int boundary = sentenceIterator.next();
 
         while (boundary != BreakIterator.DONE) {
             int sentenceEnd = sentenceIterator.current();
-            int countWordsInSentence = countWords(sentenceStart, sentenceEnd);
+            int wordsInSentenceCount = countWords(sentenceStart, sentenceEnd);
 
-            if (chunkWordCount + countWordsInSentence > maxNumberWordsPerChunk) {
+            if (chunkWordCount + wordsInSentenceCount > maxNumberWordsPerChunk) {
                 // over the max chunk size, roll back to the last sentence
 
+                int nextChunkWordCount = wordsInSentenceCount;
                 if (chunkWordCount > 0) {
                     // add a new chunk containing all the input up to this sentence
                     chunks.add(input.substring(chunkStart, chunkEnd));
-                    chunkStart = chunkEnd;
-                    chunkWordCount = countWordsInSentence; // the next chunk will contain this sentence
+
+                    if (includePrecedingSentence) {
+                        if (wordsInPrecedingSentenceCount + wordsInSentenceCount > maxNumberWordsPerChunk) {
+                            // cut the last sentence
+                            int numWordsToSkip = numWordsToSkipInPreviousSentence(wordsInPrecedingSentenceCount, maxNumberWordsPerChunk);
+
+                            chunkStart = skipWords(input, previousSentenceStart, numWordsToSkip);
+                            chunkWordCount = (wordsInPrecedingSentenceCount - numWordsToSkip) + wordsInSentenceCount;
+                        } else {
+                            chunkWordCount = wordsInPrecedingSentenceCount + wordsInSentenceCount;
+                            chunkStart = previousSentenceStart;
+                        }
+
+                        nextChunkWordCount = chunkWordCount;
+                    } else {
+                        chunkStart = chunkEnd;
+                        chunkWordCount = wordsInSentenceCount; // the next chunk will contain this sentence
+                    }
                 }
 
-                if (countWordsInSentence > maxNumberWordsPerChunk) {
-                    // This sentence is bigger than the max chunk size.
+                // Is the next chunk larger than max chunk size?
+                // If so split it
+                if (nextChunkWordCount > maxNumberWordsPerChunk) {
+                    // This sentence (and optional overlap) is bigger than the max chunk size.
                     // Split the sentence on the word boundary
                     var sentenceSplits = splitLongSentence(
                         input.substring(chunkStart, sentenceEnd),
@@ -113,7 +136,12 @@ public List<String> chunk(String input, int maxNumberWordsPerChunk) {
                     chunkWordCount = sentenceSplits.get(i).wordCount();
                 }
             } else {
-                chunkWordCount += countWordsInSentence;
+                chunkWordCount += wordsInSentenceCount;
+            }
+
+            if (includePrecedingSentence) {
+                previousSentenceStart = sentenceStart;
+                wordsInPrecedingSentenceCount = wordsInSentenceCount;
             }
 
             sentenceStart = sentenceEnd;
@@ -133,6 +161,45 @@ static List<WordBoundaryChunker.ChunkPosition> splitLongSentence(String text, in
         return new WordBoundaryChunker().chunkPositions(text, maxNumberOfWords, overlap);
     }
 
+    static int numWordsToSkipInPreviousSentence(int wordsInPrecedingSentenceCount, int maxNumberWordsPerChunk) {
+        var maxWordsInOverlap = maxWordsInOverlap(maxNumberWordsPerChunk);
+        if (wordsInPrecedingSentenceCount > maxWordsInOverlap) {
+            return wordsInPrecedingSentenceCount - maxWordsInOverlap;
+        } else {
+            return 0;
+        }
+    }
+
+    static int maxWordsInOverlap(int maxNumberWordsPerChunk) {
+        return Math.min(maxNumberWordsPerChunk / 2, 20);
+    }
+
+    private int skipWords(String input, int start, int numWords) {
+        var itr = BreakIterator.getWordInstance(Locale.ROOT);
+        itr.setText(input);
+        return skipWords(start, numWords, itr);
+    }
+
+    static int skipWords(int start, int numWords, BreakIterator wordIterator) {
+        wordIterator.preceding(start); // start of the current word
+
+        int boundary = wordIterator.current();
+        int wordCount = 0;
+        while (boundary != BreakIterator.DONE && wordCount < numWords) {
+            int wordStatus = wordIterator.getRuleStatus();
+            if (wordStatus != BreakIterator.WORD_NONE) {
+                wordCount++;
+            }
+            boundary = wordIterator.next();
+        }
+
+        if (boundary == BreakIterator.DONE) {
+            return wordIterator.last();
+        } else {
+            return boundary;
+        }
+    }
+
     private int countWords(int start, int end) {
         return countWords(start, end, this.wordIterator);
     }
@@ -157,6 +224,6 @@ static int countWords(int start, int end, BreakIterator wordIterator) {
     }
 
     private static int overlapForChunkSize(int chunkSize) {
-        return (chunkSize - 1) / 2;
+        return Math.min(20, (chunkSize - 1) / 2);
     }
 }

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

@@ -13,6 +13,7 @@
 import org.elasticsearch.common.ValidationException;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.inference.ChunkingSettings;
 import org.elasticsearch.inference.ChunkingStrategy;
 import org.elasticsearch.inference.ModelConfigurations;
@@ -30,16 +31,25 @@ public class SentenceBoundaryChunkingSettings implements ChunkingSettings {
     private static final ChunkingStrategy STRATEGY = ChunkingStrategy.SENTENCE;
     private static final Set<String> VALID_KEYS = Set.of(
         ChunkingSettingsOptions.STRATEGY.toString(),
-        ChunkingSettingsOptions.MAX_CHUNK_SIZE.toString()
+        ChunkingSettingsOptions.MAX_CHUNK_SIZE.toString(),
+        ChunkingSettingsOptions.SENTENCE_OVERLAP.toString()
     );
+
+    private static int DEFAULT_OVERLAP = 0;
+
     protected final int maxChunkSize;
+    protected int sentenceOverlap = DEFAULT_OVERLAP;
 
-    public SentenceBoundaryChunkingSettings(Integer maxChunkSize) {
+    public SentenceBoundaryChunkingSettings(Integer maxChunkSize, @Nullable Integer sentenceOverlap) {
         this.maxChunkSize = maxChunkSize;
+        this.sentenceOverlap = sentenceOverlap == null ? DEFAULT_OVERLAP : sentenceOverlap;
     }
 
     public SentenceBoundaryChunkingSettings(StreamInput in) throws IOException {
         maxChunkSize = in.readInt();
+        if (in.getTransportVersion().onOrAfter(TransportVersions.CHUNK_SENTENCE_OVERLAP_SETTING_ADDED)) {
+            sentenceOverlap = in.readVInt();
+        }
     }
 
     public static SentenceBoundaryChunkingSettings fromMap(Map<String, Object> map) {
@@ -59,11 +69,24 @@ public static SentenceBoundaryChunkingSettings fromMap(Map<String, Object> map)
             validationException
         );
 
+        Integer sentenceOverlap = ServiceUtils.extractOptionalPositiveInteger(
+            map,
+            ChunkingSettingsOptions.SENTENCE_OVERLAP.toString(),
+            ModelConfigurations.CHUNKING_SETTINGS,
+            validationException
+        );
+
+        if (sentenceOverlap != null && sentenceOverlap > 1) {
+            validationException.addValidationError(
+                ChunkingSettingsOptions.SENTENCE_OVERLAP.toString() + "[" + sentenceOverlap + "] must be either 0 or 1"
+            ); // todo better
+        }
+
         if (validationException.validationErrors().isEmpty() == false) {
             throw validationException;
         }
 
-        return new SentenceBoundaryChunkingSettings(maxChunkSize);
+        return new SentenceBoundaryChunkingSettings(maxChunkSize, sentenceOverlap);
     }
 
     @Override
@@ -72,6 +95,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         {
             builder.field(ChunkingSettingsOptions.STRATEGY.toString(), STRATEGY);
             builder.field(ChunkingSettingsOptions.MAX_CHUNK_SIZE.toString(), maxChunkSize);
+            builder.field(ChunkingSettingsOptions.SENTENCE_OVERLAP.toString(), sentenceOverlap);
         }
         builder.endObject();
         return builder;
@@ -90,6 +114,9 @@ public TransportVersion getMinimalSupportedVersion() {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeInt(maxChunkSize);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.CHUNK_SENTENCE_OVERLAP_SETTING_ADDED)) {
+            out.writeVInt(sentenceOverlap);
+        }
     }
 
     @Override

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

@@ -52,7 +52,7 @@ public static WordBoundaryChunkingSettings fromMap(Map<String, Object> map) {
         var invalidSettings = map.keySet().stream().filter(key -> VALID_KEYS.contains(key) == false).toArray();
         if (invalidSettings.length > 0) {
             validationException.addValidationError(
-                Strings.format("Sentence based chunking settings can not have the following settings: %s", Arrays.toString(invalidSettings))
+                Strings.format("Word based chunking settings can not have the following settings: %s", Arrays.toString(invalidSettings))
             );
         }
 

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

@@ -56,7 +56,7 @@ private Map<Map<String, Object>, ChunkingSettings> chunkingSettingsMapToChunking
                 ChunkingSettingsOptions.MAX_CHUNK_SIZE.toString(),
                 maxChunkSize
             ),
-            new SentenceBoundaryChunkingSettings(maxChunkSize)
+            new SentenceBoundaryChunkingSettings(maxChunkSize, 1)
         );
     }
 }

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

@@ -25,7 +25,7 @@ public static ChunkingSettings createRandomChunkingSettings() {
                 return new WordBoundaryChunkingSettings(maxChunkSize, randomIntBetween(1, maxChunkSize / 2));
             }
             case SENTENCE -> {
-                return new SentenceBoundaryChunkingSettings(randomNonNegativeInt());
+                return new SentenceBoundaryChunkingSettings(randomNonNegativeInt(), randomBoolean() ? 0 : 1);
             }
             default -> throw new IllegalArgumentException("Unsupported random strategy [" + randomStrategy + "]");
         }