Skip to content

Commit

Permalink
Adjust dashscope model batch
Browse files Browse the repository at this point in the history
Signed-off-by: junjie.jiang <[email protected]>
  • Loading branch information
junjiejiangjjj committed Dec 20, 2024
1 parent 6334590 commit 477d85e
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 4 deletions.
8 changes: 7 additions & 1 deletion internal/util/function/ali_embedding_provider.go
Original file line number Diff line number Diff line change
Expand Up @@ -95,14 +95,20 @@ func NewAliDashScopeEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functio
if err != nil {
return nil, err
}

maxBatch := 25
if modelName == TextEmbeddingV3 {
maxBatch = 6
}

provider := AliEmbeddingProvider{
client: c,
fieldDim: fieldDim,
modelName: modelName,
embedDimParam: dim,
// TextEmbedding only supports dense embedding
outputType: "dense",
maxBatch: 25,
maxBatch: maxBatch,
timeoutSec: 30,
}
return &provider, nil
Expand Down
6 changes: 3 additions & 3 deletions internal/util/function/bedrock_embedding_provider.go
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ func NewBedrockEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSche
}
var awsAccessKeyId, awsSecretAccessKey, region, modelName string
var dim int64
normalize := false
normalize := true

for _, param := range functionSchema.Params {
switch strings.ToLower(param.Key) {
Expand All @@ -110,8 +110,8 @@ func NewBedrockEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSche
region = param.Value
case normalizeParamKey:
switch strings.ToLower(param.Value) {
case "true":
normalize = true
case "false":
normalize = false
default:
return nil, fmt.Errorf("Illegal [%s:%s] param, ", normalizeParamKey, param.Value)
}
Expand Down

0 comments on commit 477d85e

Please sign in to comment.