56 lines
1.8 KiB
Go
56 lines
1.8 KiB
Go
package eino
|
||
|
||
import (
|
||
"context"
|
||
"rag/consts/model"
|
||
|
||
"github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive"
|
||
"github.com/cloudwego/eino-ext/components/document/transformer/splitter/semantic"
|
||
"github.com/cloudwego/eino/schema"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
)
|
||
|
||
// SemanticSplitDocument 语义分割文档
|
||
func SemanticSplitDocument(ctx context.Context, docs []*schema.Document, vectorModel model.ModelConfigType) (res []*schema.Document, err error) {
|
||
// 默认分隔符（支持中英文）
|
||
separators := []string{"\n\n", "\n", "。", "！", "？", "；", ".", "!", "?", ";"}
|
||
// 读取配置，使用合理的默认值
|
||
bufferSize := g.Cfg().MustGet(ctx, "eino.splitter.bufferSize").Int()
|
||
minChunkSize := g.Cfg().MustGet(ctx, "eino.splitter.minChunkSize").Int()
|
||
percentile := g.Cfg().MustGet(ctx, "eino.splitter.percentile").Float64()
|
||
batchSize := g.Cfg().MustGet(ctx, "eino.splitter.batchSize").Int()
|
||
|
||
embedder, err := GetTenantEmbedderByType(ctx, vectorModel)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
splitter, err := semantic.NewSplitter(ctx, &semantic.Config{
|
||
Embedding: NewBatchEmbedder(embedder, batchSize),
|
||
BufferSize: bufferSize,
|
||
MinChunkSize: minChunkSize,
|
||
Percentile: percentile,
|
||
Separators: separators,
|
||
})
|
||
if err != nil {
|
||
return
|
||
}
|
||
return splitter.Transform(ctx, docs)
|
||
}
|
||
|
||
// RecursiveSplitDocument 递归分割文档
|
||
func RecursiveSplitDocument(ctx context.Context, docs []*schema.Document) (res []*schema.Document, err error) {
|
||
// 默认分隔符（支持中英文）
|
||
separators := []string{"\n\n", "\n", "。", "！", "？", "；", ".", "!", "?", ";"}
|
||
splitter, err := recursive.NewSplitter(ctx, &recursive.Config{
|
||
ChunkSize: 512,
|
||
OverlapSize: 100,
|
||
KeepType: recursive.KeepTypeNone,
|
||
Separators: separators,
|
||
})
|
||
if err != nil {
|
||
return
|
||
}
|
||
return splitter.Transform(ctx, docs)
|
||
}
|