add new textspliter option: lenfunc

textsplitter/options.go

@@ -5,6 +5,7 @@ type Options struct {
 	ChunkSize         int
 	ChunkOverlap      int
 	Separators        []string
+	LenFunc           func(string) int
 	ModelName         string
 	EncodingName      string
 	AllowedSpecial    []string
@@ -20,6 +21,7 @@ func DefaultOptions() Options {
 		ChunkSize:    _defaultTokenChunkSize,
 		ChunkOverlap: _defaultTokenChunkOverlap,
 		Separators:   []string{"\n\n", "\n", " ", ""},
+		LenFunc:      defaultLenFunc,
 
 		ModelName:         _defaultTokenModelName,
 		EncodingName:      _defaultTokenEncoding,
@@ -52,6 +54,13 @@ func WithSeparators(separators []string) Option {
 	}
 }
 
+// WithLenFunc sets the lenfunc for a text splitter.
+func WithLenFunc(lenFunc func(string) int) Option {
+	return func(o *Options) {
+		o.LenFunc = lenFunc
+	}
+}
+
 // WithModelName sets the model name for a text splitter.
 func WithModelName(modelName string) Option {
 	return func(o *Options) {
@@ -107,3 +116,7 @@ func WithReferenceLinks(referenceLinks bool) Option {
 		o.ReferenceLinks = referenceLinks
 	}
 }
+
+func defaultLenFunc(s string) int {
+	return len(s)
+}

textsplitter/recursive_character.go

@@ -10,6 +10,7 @@ type RecursiveCharacter struct {
 	Separators   []string
 	ChunkSize    int
 	ChunkOverlap int
+	LenFunc      func(string) int
 }
 
 // NewRecursiveCharacter creates a new recursive character splitter with default values. By
@@ -25,6 +26,7 @@ func NewRecursiveCharacter(opts ...Option) RecursiveCharacter {
 		Separators:   options.Separators,
 		ChunkSize:    options.ChunkSize,
 		ChunkOverlap: options.ChunkOverlap,
+		LenFunc:      options.LenFunc,
 	}
 
 	return s
@@ -50,13 +52,13 @@ func (s RecursiveCharacter) SplitText(text string) ([]string, error) {
 
 	// Merge the splits, recursively splitting larger texts.
 	for _, split := range splits {
-		if len(split) < s.ChunkSize {
+		if s.LenFunc(split) < s.ChunkSize {
 			goodSplits = append(goodSplits, split)
 			continue
 		}
 
 		if len(goodSplits) > 0 {
-			mergedText := mergeSplits(goodSplits, separator, s.ChunkSize, s.ChunkOverlap)
+			mergedText := mergeSplits(goodSplits, separator, s.ChunkSize, s.ChunkOverlap, s.LenFunc)
 
 			finalChunks = append(finalChunks, mergedText...)
 			goodSplits = make([]string, 0)
@@ -74,7 +76,7 @@ func (s RecursiveCharacter) SplitText(text string) ([]string, error) {
 	}
 
 	if len(goodSplits) > 0 {
-		mergedText := mergeSplits(goodSplits, separator, s.ChunkSize, s.ChunkOverlap)
+		mergedText := mergeSplits(goodSplits, separator, s.ChunkSize, s.ChunkOverlap, s.LenFunc)
 		finalChunks = append(finalChunks, mergedText...)
 	}
 

textsplitter/recursive_character_test.go

@@ -2,6 +2,7 @@ package textsplitter
 
 import (
 	"testing"
+	"unicode/utf8"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -17,8 +18,20 @@ func TestRecursiveCharacterSplitter(t *testing.T) {
 		chunkSize    int
 		separators   []string
 		expectedDocs []schema.Document
+		lenFunc      func(string) int
 	}
 	testCases := []testCase{
+		{
+			text:         "哈里森\n很高兴遇见你\n欢迎你来中国",
+			chunkOverlap: 0,
+			chunkSize:    10,
+			separators:   []string{"\n\n", "\n", " "},
+			lenFunc:      utf8.RuneCountInString,
+			expectedDocs: []schema.Document{
+				{PageContent: "哈里森\n很高兴遇见你", Metadata: map[string]any{}},
+				{PageContent: "欢迎你来中国", Metadata: map[string]any{}},
+			},
+		},
 		{
 			text:         "Hi, Harrison. \nI am glad to meet you",
 			chunkOverlap: 1,
@@ -102,6 +115,9 @@ Bye!
 		splitter.ChunkOverlap = tc.chunkOverlap
 		splitter.ChunkSize = tc.chunkSize
 		splitter.Separators = tc.separators
+		if tc.lenFunc != nil {
+			splitter.LenFunc = tc.lenFunc
+		}
 
 		docs, err := CreateDocuments(splitter, []string{tc.text}, nil)
 		require.NoError(t, err)

textsplitter/split_documents.go

@@ -68,15 +68,15 @@ func joinDocs(docs []string, separator string) string {
 }
 
 // mergeSplits merges smaller splits into splits that are closer to the chunkSize.
-func mergeSplits(splits []string, separator string, chunkSize int, chunkOverlap int) []string { //nolint:cyclop
+func mergeSplits(splits []string, separator string, chunkSize int, chunkOverlap int, lenFunc func(string) int) []string { //nolint:cyclop
 	docs := make([]string, 0)
 	currentDoc := make([]string, 0)
 	total := 0
 
 	for _, split := range splits {
-		totalWithSplit := total + len(split)
+		totalWithSplit := total + lenFunc(split)
 		if len(currentDoc) != 0 {
-			totalWithSplit += len(separator)
+			totalWithSplit += lenFunc(separator)
 		}
 
 		maybePrintWarning(total, chunkSize)
@@ -86,19 +86,19 @@ func mergeSplits(splits []string, separator string, chunkSize int, chunkOverlap
 				docs = append(docs, doc)
 			}
 
-			for shouldPop(chunkOverlap, chunkSize, total, len(split), len(separator), len(currentDoc)) {
-				total -= len(currentDoc[0]) //nolint:gosec
+			for shouldPop(chunkOverlap, chunkSize, total, lenFunc(split), lenFunc(separator), len(currentDoc)) {
+				total -= lenFunc(currentDoc[0]) //nolint:gosec
 				if len(currentDoc) > 1 {
-					total -= len(separator)
+					total -= lenFunc(separator)
 				}
 				currentDoc = currentDoc[1:] //nolint:gosec
 			}
 		}
 
 		currentDoc = append(currentDoc, split)
-		total += len(split)
+		total += lenFunc(split)
 		if len(currentDoc) > 1 {
-			total += len(separator)
+			total += lenFunc(separator)
 		}
 	}