Update TextSplitter to generate TextChunk

Author: intitni

Tool/Sources/LangChain/DocumentTransformer/RecursiveCharacterTextSplitter.swift

@@ -19,7 +19,7 @@ public class RecursiveCharacterTextSplitter: TextSplitter {
     ///    - chunkOverlap: The maximum overlap between chunks.
     ///    - lengthFunction: A function to compute the length of text.
     public init(
-        separators: [String] = ["\n\n", "\r\n", "\n", "\r", " ", ""],
+        separators: [String],
         chunkSize: Int = 4000,
         chunkOverlap: Int = 200,
         lengthFunction: @escaping (String) -> Int = { $0.count }
@@ -39,7 +39,7 @@ public class RecursiveCharacterTextSplitter: TextSplitter {
     ///    - chunkOverlap: The maximum overlap between chunks.
     ///    - lengthFunction: A function to compute the length of text.
     public init(
-        separatorSet: TextSplitterSeparatorSet,
+        separatorSet: TextSplitterSeparatorSet = .default,
         chunkSize: Int = 4000,
         chunkOverlap: Int = 200,
         lengthFunction: @escaping (String) -> Int = { $0.count }
@@ -51,12 +51,12 @@ public class RecursiveCharacterTextSplitter: TextSplitter {
         separators = separatorSet.separators
     }
 
-    public func split(text: String) async throws -> [String] {
-        return split(text: text, separators: separators)
+    public func split(text: String) async throws -> [TextChunk] {
+        return split(text: text, separators: separators, startIndex: 0)
     }
 
-    private func split(text: String, separators: [String]) -> [String] {
-        var finalChunks = [String]()
+    private func split(text: String, separators: [String], startIndex: Int) -> [TextChunk] {
+        var finalChunks = [TextChunk]()
 
         // Get appropriate separator to use
         let firstSeparatorIndex = separators.firstIndex {
@@ -83,12 +83,12 @@ public class RecursiveCharacterTextSplitter: TextSplitter {
             nextSeparators = []
         }
 
-        let splits = split(text: text, separator: separator)
+        let splits = split(text: text, separator: separator, startIndex: startIndex)
 
         // Now go merging things, recursively splitting longer texts.
-        var goodSplits = [String]()
+        var goodSplits = [TextChunk]()
         for s in splits {
-            if lengthFunction(s) < chunkSize {
+            if lengthFunction(s.text) < chunkSize {
                 goodSplits.append(s)
             } else {
                 if !goodSplits.isEmpty {
@@ -99,7 +99,11 @@ public class RecursiveCharacterTextSplitter: TextSplitter {
                 if nextSeparators.isEmpty {
                     finalChunks.append(s)
                 } else {
-                    let other_info = split(text: s, separators: nextSeparators)
+                    let other_info = split(
+                        text: s.text,
+                        separators: nextSeparators,
+                        startIndex: s.startUTF16Offset
+                    )
                     finalChunks.append(contentsOf: other_info)
                 }
             }

Tool/Sources/LangChain/DocumentTransformer/TextSplitter.swift

@@ -11,7 +11,7 @@ public protocol TextSplitter: DocumentTransformer {
     var lengthFunction: (String) -> Int { get }
 
     /// Split text into multiple components.
-    func split(text: String) async throws -> [String]
+    func split(text: String) async throws -> [TextChunk]
 }
 
 public extension TextSplitter {
@@ -26,7 +26,7 @@ public extension TextSplitter {
         for (text, metadata) in zip(texts, metadata) {
             let chunks = try await split(text: text)
             for chunk in chunks {
-                let document = Document(pageContent: chunk, metadata: metadata)
+                let document = Document(pageContent: chunk.text, metadata: metadata)
                 documents.append(document)
             }
         }
@@ -50,31 +50,48 @@ public extension TextSplitter {
     }
 }
 
+public struct TextChunk {
+    public var text: String
+    public var startUTF16Offset: Int
+    public var endUTF16Offset: Int
+}
+
 public extension TextSplitter {
     /// Merge small splits to just fit in the chunk size.
-    func mergeSplits(_ splits: [String]) -> [String] {
+    func mergeSplits(_ splits: [TextChunk]) -> [TextChunk] {
         let chunkOverlap = chunkOverlap < chunkSize ? chunkOverlap : 0
 
-        var chunks = [String]()
-        var currentChunk = [String]()
-        var overlappingChunks = [String]()
+        var chunks = [TextChunk]()
+        var currentChunk = [TextChunk]()
+        var overlappingChunks = [TextChunk]()
         var currentChunkSize = 0
-        
-        func join(_ a: [String], _ b: [String]) -> String {
-            return (a + b).joined().trimmingCharacters(in: .whitespaces)
+
+        func join(_ a: [TextChunk], _ b: [TextChunk]) -> TextChunk? {
+            let text = (a + b).map(\.text).joined()
+            var l = Int.max
+            var u = 0
+            
+            for chunk in a + b {
+                l = min(l, chunk.startUTF16Offset)
+                u = max(u, chunk.endUTF16Offset)
+            }
+            
+            guard l < u else { return nil }
+            
+            return .init(text: text, startUTF16Offset: l, endUTF16Offset: u)
         }
 
-        for text in splits {
-            let textLength = lengthFunction(text)
+        for chunk in splits {
+            let textLength = lengthFunction(chunk.text)
             if currentChunkSize + textLength > chunkSize {
-                let currentChunkText = join(overlappingChunks, currentChunk)
+                guard let currentChunkText = join(overlappingChunks, currentChunk) else { continue }
                 chunks.append(currentChunkText)
 
                 overlappingChunks = []
                 var overlappingSize = 0
                 // use small chunks as overlap if possible
                 for chunk in currentChunk.reversed() {
-                    let length = lengthFunction(chunk)
+                    let length = lengthFunction(chunk.text)
                     if overlappingSize + length > chunkOverlap { break }
                     if overlappingSize + length + textLength > chunkSize { break }
                     overlappingSize += length
@@ -90,46 +107,57 @@ public extension TextSplitter {
 //                }
 
                 currentChunkSize = overlappingSize + textLength
-                currentChunk = [text]
+                currentChunk = [chunk]
             } else {
                 currentChunkSize += textLength
-                currentChunk.append(text)
+                currentChunk.append(chunk)
             }
         }
 
-        if !currentChunk.isEmpty {
-            chunks.append(join(overlappingChunks, currentChunk))
+        if !currentChunk.isEmpty, let joinedChunks = join(overlappingChunks, currentChunk) {
+            chunks.append(joinedChunks)
+        } else {
+            chunks.append(contentsOf: overlappingChunks)
+            chunks.append(contentsOf: currentChunk)
         }
 
         return chunks
     }
 
     /// Split the text by separator.
-    func split(text: String, separator: String) -> [String] {
-        guard !separator.isEmpty else {
-            return [text]
-        }
-
+    func split(text: String, separator: String, startIndex: Int = 0) -> [TextChunk] {
         let pattern = "(\(separator))"
-        if let regex = try? NSRegularExpression(pattern: pattern) {
+        if !separator.isEmpty, let regex = try? NSRegularExpression(pattern: pattern) {
             let matches = regex.matches(in: text, range: NSRange(text.startIndex..., in: text))
-            var all = [String]()
+            var all = [TextChunk]()
             var start = text.startIndex
             for match in matches {
                 guard let range = Range(match.range, in: text) else { break }
                 guard range.lowerBound > start else { break }
                 let result = text[start..<range.lowerBound]
-                start = range.lowerBound
                 if !result.isEmpty {
-                    all.append(String(result))
+                    all.append(.init(
+                        text: String(result),
+                        startUTF16Offset: start.utf16Offset(in: text) + startIndex,
+                        endUTF16Offset: range.lowerBound.utf16Offset(in: text) + startIndex
+                    ))
                 }
+                start = range.lowerBound
             }
             if start < text.endIndex {
-                all.append(String(text[start...]))
+                all.append(.init(
+                    text: String(text[start...]),
+                    startUTF16Offset: start.utf16Offset(in: text) + startIndex,
+                    endUTF16Offset: text.endIndex.utf16Offset(in: text) + startIndex
+                ))
             }
             return all
         } else {
-            return [text]
+            return [.init(
+                text: text,
+                startUTF16Offset: text.startIndex.utf16Offset(in: text) + startIndex,
+                endUTF16Offset: text.endIndex.utf16Offset(in: text) + startIndex
+            )]
         }
     }
 }

Tool/Sources/LangChain/DocumentTransformer/TextSplitterSeparatorSet.swift

@@ -347,5 +347,9 @@ public struct TextSplitterSeparatorSet: ExpressibleByArrayLiteral {
             "",
         ]
     }
+
+    public static var `default`: TextSplitterSeparatorSet {
+        ["\n\n", "\r\n", "\n", "\r", " ", ""]
+    }
 }