ben/coze-studio
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: support EMBEDDING_MAX_BATCH_SIZE (#311)

Browse Source
This commit is contained in:
N3ko
2025-07-30 15:31:54 +08:00
committed by GitHub
parent f93f26fc48
commit bb74272385
9 changed files with 85 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
backend/infra/impl/embedding/wrap/ollama.go
View File

@@ -24,10 +24,10 @@ import (
contract "github.com/coze-dev/coze-studio/backend/infra/contract/embedding"
)
func NewOllamaEmbedder(ctx context.Context, config *ollama.EmbeddingConfig, dimensions int64) (contract.Embedder, error) {
func NewOllamaEmbedder(ctx context.Context, config *ollama.EmbeddingConfig, dimensions int64, batchSize int) (contract.Embedder, error) {
emb, err := ollama.NewEmbedder(ctx, config)
if err != nil {
return nil, err
}
return &denseOnlyWrap{dims: dimensions, Embedder: emb}, nil
return &denseOnlyWrap{dims: dimensions, batchSize: batchSize, Embedder: emb}, nil
}