Add custom splitter example

Author: johnoliver

samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/rag/DocumentSplittingExample.java

@@ -4,6 +4,12 @@
 import com.microsoft.semantic.kernel.rag.splitting.Chunk;
 import com.microsoft.semantic.kernel.rag.splitting.Document;
 import com.microsoft.semantic.kernel.rag.splitting.Splitter;
+import com.microsoft.semantic.kernel.rag.splitting.TextSplitter;
+import com.microsoft.semantic.kernel.rag.splitting.document.TextDocument;
+import com.microsoft.semantic.kernel.rag.splitting.overlap.NoOverlapCondition;
+import com.microsoft.semantic.kernel.rag.splitting.splitconditions.CountSplitCondition;
+import com.microsoft.semantic.kernel.rag.splitting.splitconditions.SplitPoint;
+import com.microsoft.semantickernel.implementation.EmbeddedResourceLoader;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.net.URI;
@@ -12,11 +18,14 @@
 import java.net.http.HttpResponse;
 import java.net.http.HttpResponse.BodyHandlers;
 import java.util.List;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 import org.apache.pdfbox.io.RandomAccessReadBuffer;
 import org.apache.pdfbox.pdfparser.PDFParser;
 import org.apache.pdfbox.pdmodel.PDDocument;
 import org.apache.pdfbox.text.PDFTextStripper;
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 public class DocumentSplittingExample {
 
@@ -46,6 +55,11 @@ public Flux<String> getContent() {
     }
 
     public static void main(String[] args) throws IOException, InterruptedException {
+        useCustomChunker();
+        useInbuiltChunker();
+    }
+
+    private static void useInbuiltChunker() throws IOException, InterruptedException {
         byte[] pdfBytes = getPdfDoc();
         PDFDocument pdfDoc = new PDFDocument(pdfBytes);
 
@@ -68,6 +82,70 @@ public static void main(String[] args) throws IOException, InterruptedException
             });
     }
 
+    public static void useCustomChunker() throws IOException, InterruptedException {
+
+        String example = EmbeddedResourceLoader.readFile("example.md",
+            DocumentSplittingExample.class);
+
+        // Define how we are splitting tokens, in this case we are splitting on headers of an md file
+        // i.e <new line> followed by one or more # characters
+        TextSplitter textSplitter = (doc, numTokens) -> {
+            // Split on headers
+            Pattern pattern = Pattern.compile("(\\r?\\n|\\r)\s*#+", Pattern.MULTILINE);
+
+            Flux<Integer> splitPoints = Flux.fromStream(pattern.matcher(doc).results())
+                .map(window -> window.start());
+
+            return createWindows(doc, splitPoints);
+        };
+
+        // Split into single sections
+        CountSplitCondition condition = new CountSplitCondition(1, textSplitter);
+
+        Splitter splitter = Splitter
+            .builder()
+            .addChunkEndCondition(condition)
+            // No overlap
+            .setOverlapCondition(NoOverlapCondition.build())
+            // Tidy up the text
+            .trimWhitespace()
+            .build();
+
+        String chunks = splitter
+            .splitDocument(new TextDocument(example))
+            .collectList()
+            .map(it -> it.stream()
+                .map(chunk -> chunk.getContents())
+                .collect(Collectors.joining("\n============\n")))
+            .block();
+
+        System.out.println(chunks);
+    }
+
+    /*
+     * Transforms: [ 2, 10, 20, 100 ] -> [ (0, 2), (2, 10), (10, 20), (20, 100), (100, <doc length>)
+     * ]
+     */
+    private static List<SplitPoint> createWindows(String doc, Flux<Integer> splitPoints) {
+        return Flux.concat(
+            Flux.just(0),
+            splitPoints,
+            Flux.just(doc.length()))
+            .window(2, 1)
+            .concatMap(window -> {
+                return window.collectList()
+                    .flatMap(list -> {
+                        if (list.size() <= 1) {
+                            return Mono.empty();
+                        }
+                        return Mono.just(
+                            new SplitPoint(list.get(0), list.get(1)));
+                    });
+            })
+            .collectList()
+            .block();
+    }
+
     private static byte[] getPdfDoc() throws IOException, InterruptedException {
         HttpResponse<byte[]> doc = HttpClient.newHttpClient()
             .send(HttpRequest.newBuilder()

samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/resources/com/microsoft/semantickernel/samples/syntaxexamples/rag/example.md

@@ -0,0 +1,22 @@
+## Section 1
+
+Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna
+aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis
+aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint
+occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
+
+## Section 2
+
+Another section.
+
+### Subsection 1
+
+1, 2, 3, 4, 5, 6, 7, 8, 9, 10.
+
+# Section 3
+
+This is the last section.
+
+```
+some code
+```

samples/semantickernel-sample-plugins/semantickernel-text-splitter-plugin/src/main/java/com/microsoft/semantic/kernel/rag/splitting/Splitter.java

@@ -118,15 +118,16 @@ private static List<Chunk> chunkDocument(List<ChunkEndCondition> chunkEndConditi
                         "This entier chunk consists of overlapped data, this will result in infinite loop. Skipping this chunk.");
 
                     // previous chunk should already contain this text..skip it
-                    doc = doc.substring(previousChunkEndIndex, doc.length());
+                    doc = doc.substring(Math.min(previousChunkEndIndex, doc.length()),
+                        doc.length());
 
-                    previousChunkEndIndex = 0;
+                    previousChunkEndIndex = -1;
                     continue;
                 }
 
                 int overlapIndex = overlapCondition.getOverlapIndex(chunkText);
                 previousChunkEndIndex = chunkText.length() - overlapIndex;
-                doc = doc.substring(overlapIndex, doc.length());
+                doc = doc.substring(Math.min(overlapIndex, doc.length()), doc.length());
 
                 chunks.add(new Chunk(chunkText));
             } else {

samples/semantickernel-sample-plugins/semantickernel-text-splitter-plugin/src/main/java/com/microsoft/semantic/kernel/rag/splitting/TextSplitter.java

@@ -1,8 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 package com.microsoft.semantic.kernel.rag.splitting;
 
-import com.microsoft.semantic.kernel.rag.splitting.splitconditions.SplitPoints;
-
+import com.microsoft.semantic.kernel.rag.splitting.splitconditions.SplitPoint;
 import java.util.List;
 
 /**
@@ -16,7 +15,9 @@ public interface TextSplitter {
      * @param doc the document to split
      * @return the split points
      */
-    List<SplitPoints> getSplitPoints(String doc);
+    default List<SplitPoint> getSplitPoints(String doc) {
+        return getNSplitPoints(doc, Integer.MAX_VALUE);
+    }
 
     /**
      * Get the first n split points for the given document.
@@ -25,5 +26,5 @@ public interface TextSplitter {
      * @param n   the number of split points to get
      * @return the split points
      */
-    List<SplitPoints> getNSplitPoints(String doc, int n);
+    List<SplitPoint> getNSplitPoints(String doc, int n);
 }

samples/semantickernel-sample-plugins/semantickernel-text-splitter-plugin/src/main/java/com/microsoft/semantic/kernel/rag/splitting/overlap/CountOverlapCondition.java

@@ -3,7 +3,7 @@
 
 import com.microsoft.semantic.kernel.rag.splitting.OverlapCondition;
 import com.microsoft.semantic.kernel.rag.splitting.TextSplitter;
-import com.microsoft.semantic.kernel.rag.splitting.splitconditions.SplitPoints;
+import com.microsoft.semantic.kernel.rag.splitting.splitconditions.SplitPoint;
 import java.util.List;
 
 /**
@@ -25,7 +25,7 @@ public CountOverlapCondition(int count, TextSplitter splitter) {
 
     @Override
     public int getOverlapIndex(String chunk) {
-        List<SplitPoints> splitPoints = splitter.getSplitPoints(chunk);
+        List<SplitPoint> splitPoints = splitter.getSplitPoints(chunk);
 
         if (splitPoints.size() == 0) {
             return 0;

samples/semantickernel-sample-plugins/semantickernel-text-splitter-plugin/src/main/java/com/microsoft/semantic/kernel/rag/splitting/overlap/NoOverlapCondition.java

@@ -11,6 +11,10 @@ public class NoOverlapCondition implements OverlapCondition {
     public NoOverlapCondition() {
     }
 
+    public static OverlapCondition build() {
+        return new NoOverlapCondition();
+    }
+
     @Override
     public int getOverlapIndex(String chunk) {
         return chunk.length();

samples/semantickernel-sample-plugins/semantickernel-text-splitter-plugin/src/main/java/com/microsoft/semantic/kernel/rag/splitting/overlap/PercentageOverlapCondition.java

@@ -3,7 +3,7 @@
 
 import com.microsoft.semantic.kernel.rag.splitting.OverlapCondition;
 import com.microsoft.semantic.kernel.rag.splitting.TextSplitter;
-import com.microsoft.semantic.kernel.rag.splitting.splitconditions.SplitPoints;
+import com.microsoft.semantic.kernel.rag.splitting.splitconditions.SplitPoint;
 import java.util.List;
 import org.slf4j.Logger;
 
@@ -31,11 +31,11 @@ public PercentageOverlapCondition(float percentage, TextSplitter splitter) {
 
     @Override
     public int getOverlapIndex(String chunk) {
-        List<SplitPoints> splitPoints = splitter.getSplitPoints(chunk);
+        List<SplitPoint> splitPoints = splitter.getSplitPoints(chunk);
 
         float index = chunk.length() * (100.0f - percentage) / 100.0f;
 
-        for (SplitPoints splitPoint : splitPoints) {
+        for (SplitPoint splitPoint : splitPoints) {
             if (splitPoint.getEnd() > index) {
                 return splitPoint.getStart();
             }

samples/semantickernel-sample-plugins/semantickernel-text-splitter-plugin/src/main/java/com/microsoft/semantic/kernel/rag/splitting/splitconditions/CountSplitCondition.java

@@ -21,7 +21,13 @@ public CountSplitCondition(int count, TextSplitter splitter) {
 
     @Override
     public int getEndOfNextChunk(String doc) {
-        List<SplitPoints> splitPoints = splitter.getNSplitPoints(doc, count);
+        List<SplitPoint> splitPoints = splitter.getNSplitPoints(doc, count)
+            .stream()
+            .filter(it -> it != null)
+            .filter(it -> it.getEnd() != 0)
+            .filter(it -> it.getEnd() != it.getStart())
+            .filter(it -> it.getStart() != doc.length())
+            .toList();
 
         if (splitPoints.size() < count) {
             return splitPoints.get(splitPoints.size() - 1).getEnd();

samples/semantickernel-sample-plugins/semantickernel-text-splitter-plugin/src/main/java/com/microsoft/semantic/kernel/rag/splitting/splitconditions/RegexSplitter.java

@@ -44,18 +44,13 @@ public RegexSplitter(Pattern pattern, int trivialSplitLength) {
     }
 
     @Override
-    public List<SplitPoints> getSplitPoints(String doc) {
-        return getNSplitPoints(doc, Integer.MAX_VALUE);
-    }
-
-    @Override
-    public List<SplitPoints> getNSplitPoints(String doc, int n) {
+    public List<SplitPoint> getNSplitPoints(String doc, int n) {
         Matcher matcher = pattern.matcher(doc);
 
         List<MatchResult> points = matcher.results()
             .collect(Collectors.toList());
 
-        List<SplitPoints> result = new ArrayList<>();
+        List<SplitPoint> result = new ArrayList<>();
 
         int previousEnd = 0;
         for (MatchResult point : points) {
@@ -66,19 +61,19 @@ public List<SplitPoints> getNSplitPoints(String doc, int n) {
                 trivialSplitLength)) {
                 continue;
             }
-            result.add(new SplitPoints(previousEnd, point.end()));
+            result.add(new SplitPoint(previousEnd, point.end()));
             previousEnd = point.end();
             if (result.size() >= n) {
                 break;
             }
         }
 
         if (result.size() < n && !isTrivialSplit(previousEnd, doc.length(), doc, 1)) {
-            result.add(new SplitPoints(previousEnd, doc.length()));
+            result.add(new SplitPoint(previousEnd, doc.length()));
         }
 
         if (result.isEmpty()) {
-            return List.of(new SplitPoints(0, doc.length()));
+            return List.of(new SplitPoint(0, doc.length()));
         }
 
         return result;

samples/semantickernel-sample-plugins/semantickernel-text-splitter-plugin/src/main/java/com/microsoft/semantic/kernel/rag/splitting/splitconditions/SplitPoints.java renamed to samples/semantickernel-sample-plugins/semantickernel-text-splitter-plugin/src/main/java/com/microsoft/semantic/kernel/rag/splitting/splitconditions/SplitPoint.java

@@ -5,12 +5,12 @@
  * A class that represents the start and end points of a split. I.e if splitting by word, these
  * would be the indices of the first and last char in the word within the chunk.
  */
-public class SplitPoints {
+public class SplitPoint {
 
     private final int start;
     private final int end;
 
-    public SplitPoints(int start, int end) {
+    public SplitPoint(int start, int end) {
         this.start = start;
         this.end = end;
     }