Repository: tmc/langchaingo
Pull request #163 (2023-07-02 02:07:43)

add new token splitter
### PR Checklist

- [x] Read the [Contributing documentation](https://github.com/tmc/langchaingo/blob/main/CONTRIBUTING.md).
- [x] Read the [Code of conduct documentation](https://github.com/tmc/langchaingo/blob/main/CODE_OF_CONDUCT.md).
- [x] Name your Pull Request title clearly, concisely, and prefixed with the name of the primarily affected package you changed according to [Good commit messages](https://go.dev/doc/contribute#commit_messages) (such as `memory: add interfaces for X, Y` or `util: add whizzbang helpers`).
- [x] Check that there isn't already a PR that solves the problem the same way to avoid creating a duplicate.
- [x] Provide a description in this PR that addresses **what** the PR is solving, or reference the issue that it solves (e.g. `Fixes #123`).
- [x] Describes the source of new concepts.
- [x] References existing implementations as appropriate.
- [x] Contains test coverage for new functions.
- [ ] Passes all [`golangci-lint`](https://golangci-lint.run/) checks.



Patch:
diff --git a/textsplitter/token_splitter.go b/textsplitter/token_splitter.go
new file mode 100644
index 00000000..fd9d98c0
--- /dev/null
+++ b/textsplitter/token_splitter.go
@@ -0,0 +1,75 @@
+package textsplitter
+
+import (
+	"fmt"
+
+	"github.com/pkoukk/tiktoken-go"
+)
+
+const (
+	// nolint:gosec
+	_defaultTokenModelName    = "gpt-3.5-turbo"
+	_defaultTokenEncoding     = "cl100k_base"
+	_defaultTokenChunkSize    = 512
+	_defaultTokenChunkOverlap = 100
+)
+
+// TokenSplitter is a text splitter that will split texts by tokens.
+type TokenSplitter struct {
+	ChunkSize         int
+	ChunkOverlap      int
+	ModelName         string
+	EncodingName      string
+	AllowedSpecial    []string
+	DisallowedSpecial []string
+}
+
+func NewTokenSplitter() TokenSplitter {
+	return TokenSplitter{
+		ChunkSize:         _defaultTokenChunkSize,
+		ChunkOverlap:      _defaultTokenChunkOverlap,
+		ModelName:         _defaultTokenModelName,
+		EncodingName:      _defaultTokenEncoding,
+		AllowedSpecial:    []string{},
+		DisallowedSpecial: []string{"all"},
+	}
+}
+
+// SplitText splits a text into multiple text.
+func (s TokenSplitter) SplitText(text string) ([]string, error) {
+	// Get the tokenizer
+	var tk *tiktoken.Tiktoken
+	var err error
+	if s.EncodingName != "" {
+		tk, err = tiktoken.GetEncoding(s.EncodingName)
+	} else {
+		tk, err = tiktoken.EncodingForModel(s.ModelName)
+	}
+	if err != nil {
+		return nil, fmt.Errorf("tiktoken.GetEncoding: %w", err)
+	}
+	texts := s.splitText(text, tk)
+
+	return texts, nil
+}
+
+func (s TokenSplitter) splitText(text string, tk *tiktoken.Tiktoken) []string {
+	splits := make([]string, 0)
+	inputIds := tk.Encode(text, s.AllowedSpecial, s.DisallowedSpecial)
+
+	startIdx := 0
+	curIdx := len(inputIds)
+	if startIdx+s.ChunkSize < curIdx {
+		curIdx = startIdx + s.ChunkSize
+	}
+	for startIdx < len(inputIds) {
+		chunkIds := inputIds[startIdx:curIdx]
+		splits = append(splits, tk.Decode(chunkIds))
+		startIdx += s.ChunkSize - s.ChunkOverlap
+		curIdx = startIdx + s.ChunkSize
+		if curIdx > len(inputIds) {
+			curIdx = len(inputIds)
+		}
+	}
+	return splits
+}


Test patch:
diff --git a/textsplitter/token_splitter_test.go b/textsplitter/token_splitter_test.go
new file mode 100644
index 00000000..462fdfde
--- /dev/null
+++ b/textsplitter/token_splitter_test.go
@@ -0,0 +1,87 @@
+package textsplitter
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tmc/langchaingo/schema"
+)
+
+func TestTokenSplitter(t *testing.T) {
+	t.Parallel()
+	type testCase struct {
+		text         string
+		chunkOverlap int
+		chunkSize    int
+		expectedDocs []schema.Document
+	}
+	//nolint:dupword
+	testCases := []testCase{
+		{
+			text:         "Hi.\nI'm Harrison.\n\nHow?\na\nb",
+			chunkOverlap: 1,
+			chunkSize:    20,
+			expectedDocs: []schema.Document{
+				{PageContent: "Hi.\nI'm Harrison.\n\nHow?\na\nb", Metadata: map[string]any{}},
+			},
+		},
+		{
+			text:         "Hi.\nI'm Harrison.\n\nHow?\na\nbHi.\nI'm Harrison.\n\nHow?\na\nb",
+			chunkOverlap: 1,
+			chunkSize:    40,
+			expectedDocs: []schema.Document{
+				{PageContent: "Hi.\nI'm Harrison.\n\nHow?\na\nbHi.\nI'm Harrison.\n\nHow?\na\nb", Metadata: map[string]any{}},
+			},
+		},
+		{
+			text:         "name: Harrison\nage: 30",
+			chunkOverlap: 1,
+			chunkSize:    40,
+			expectedDocs: []schema.Document{
+				{PageContent: "name: Harrison\nage: 30", Metadata: map[string]any{}},
+			},
+		},
+		{
+			text: `name: Harrison
+age: 30
+
+name: Joe
+age: 32`,
+			chunkOverlap: 1,
+			chunkSize:    40,
+			expectedDocs: []schema.Document{
+				{PageContent: "name: Harrison\nage: 30\n\nname: Joe\nage: 32", Metadata: map[string]any{}},
+			},
+		},
+		{
+			text: `Hi.
+I'm Harrison.
+
+How? Are? You?
+Okay then f f f f.
+This is a weird text to write, but gotta test the splittingggg some how.
+
+Bye!
+
+-H.`,
+			chunkOverlap: 1,
+			chunkSize:    10,
+			expectedDocs: []schema.Document{
+				{PageContent: "Hi.\nI'm Harrison.\n\nHow? Are?", Metadata: map[string]any{}},
+				{PageContent: "? You?\nOkay then f f f f.\n", Metadata: map[string]any{}},
+				{PageContent: ".\nThis is a weird text to write, but", Metadata: map[string]any{}},
+				{PageContent: " but gotta test the splittingggg some how.\n\n", Metadata: map[string]any{}},
+				{PageContent: ".\n\nBye!\n\n-H.", Metadata: map[string]any{}},
+			},
+		},
+	}
+	splitter := NewTokenSplitter()
+	for _, tc := range testCases {
+		splitter.ChunkOverlap = tc.chunkOverlap
+		splitter.ChunkSize = tc.chunkSize
+
+		docs, err := CreateDocuments(splitter, []string{tc.text}, nil)
+		assert.NoError(t, err)
+		assert.Equal(t, tc.expectedDocs, docs)
+	}
+}