Add --length flag to control model output "maxTokens"
This commit is contained in:
parent
681b52a55c
commit
8bc8312154
@ -9,10 +9,16 @@ import (
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// TODO: allow setting with flag
|
||||
const MAX_TOKENS = 256
|
||||
var (
|
||||
maxTokens int
|
||||
)
|
||||
|
||||
func init() {
|
||||
inputCmds := []*cobra.Command{newCmd, promptCmd, replyCmd}
|
||||
for _, cmd := range inputCmds {
|
||||
cmd.Flags().IntVar(&maxTokens, "length", config.OpenAI.DefaultMaxLength, "Max response length in tokens")
|
||||
}
|
||||
|
||||
rootCmd.AddCommand(
|
||||
lsCmd,
|
||||
newCmd,
|
||||
@ -268,7 +274,7 @@ var replyCmd = &cobra.Command{
|
||||
response <- HandleDelayedResponse(receiver)
|
||||
}()
|
||||
|
||||
err = CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, MAX_TOKENS, receiver)
|
||||
err = CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, maxTokens, receiver)
|
||||
if err != nil {
|
||||
Fatal("%v\n", err)
|
||||
}
|
||||
@ -348,7 +354,7 @@ var newCmd = &cobra.Command{
|
||||
response <- HandleDelayedResponse(receiver)
|
||||
}()
|
||||
|
||||
err = CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, MAX_TOKENS, receiver)
|
||||
err = CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, maxTokens, receiver)
|
||||
if err != nil {
|
||||
Fatal("%v\n", err)
|
||||
}
|
||||
@ -397,7 +403,7 @@ var promptCmd = &cobra.Command{
|
||||
|
||||
receiver := make(chan string)
|
||||
go HandleDelayedResponse(receiver)
|
||||
err := CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, MAX_TOKENS, receiver)
|
||||
err := CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, maxTokens, receiver)
|
||||
if err != nil {
|
||||
Fatal("%v\n", err)
|
||||
}
|
||||
|
@ -12,6 +12,7 @@ type Config struct {
|
||||
OpenAI struct {
|
||||
APIKey string `yaml:"apiKey"`
|
||||
DefaultModel string `yaml:"defaultModel"`
|
||||
DefaultMaxLength int `yaml:"defaultMaxLength"`
|
||||
} `yaml:"openai"`
|
||||
}
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user