File tree 1 file changed +16
-12
lines changed
1 file changed +16
-12
lines changed Original file line number Diff line number Diff line change 1
1
package ai
2
2
3
3
import (
4
- "os"
5
- "fmt"
6
4
"context"
7
5
"errors"
6
+ "fmt"
7
+ "os"
8
8
9
9
wx "github.com/IBM/watsonx-go/pkg/models"
10
10
)
@@ -14,28 +14,33 @@ const watsonxAIClientName = "watsonxai"
14
14
type WatsonxAIClient struct {
15
15
nopCloser
16
16
17
- client * wx.Client
18
- model string
19
- temperature float32
20
- topP float32
21
- topK int32
22
- maxNewTokens int
17
+ client * wx.Client
18
+ model string
19
+ temperature float32
20
+ topP float32
21
+ topK int32
22
+ maxNewTokens int
23
23
}
24
24
25
25
const (
26
26
modelMetallama = "ibm/granite-13b-chat-v2"
27
+ maxTokens = 2048
27
28
)
28
29
29
30
func (c * WatsonxAIClient ) Configure (config IAIConfig ) error {
30
- if (config .GetModel () == "" ) {
31
+ if config .GetModel () == "" {
32
+ c .model = modelMetallama
33
+ } else {
31
34
c .model = config .GetModel ()
35
+ }
36
+ if config .GetMaxTokens () == 0 {
37
+ c .maxNewTokens = maxTokens
32
38
} else {
33
- c .model = modelMetallama
39
+ c .maxNewTokens = config . GetMaxTokens ()
34
40
}
35
41
c .temperature = config .GetTemperature ()
36
42
c .topP = config .GetTopP ()
37
43
c .topK = config .GetTopK ()
38
- c .maxNewTokens = config .GetMaxTokens ()
39
44
40
45
// WatsonxAPIKeyEnvVarName = "WATSONX_API_KEY"
41
46
// WatsonxProjectIDEnvVarName = "WATSONX_PROJECT_ID"
@@ -75,7 +80,6 @@ func (c *WatsonxAIClient) GetCompletion(ctx context.Context, prompt string) (str
75
80
if result .Text == "" {
76
81
return "" , errors .New ("Expected a result, but got an empty string" )
77
82
}
78
-
79
83
return result .Text , nil
80
84
}
81
85
You can’t perform that action at this time.
0 commit comments