xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-09-18 00:16:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

chore: update InitTokenEncoders (#255)

Browse Source
This commit is contained in:
CaIon 2024-05-15 16:32:00 +08:00
parent 7b5830522a
commit a3b3e6cc38
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
service/token_counter.go
View File

@ -26,6 +26,7 @@ func InitTokenEncoders() {
} }
defaultTokenEncoder = gpt35TokenEncoder defaultTokenEncoder = gpt35TokenEncoder
gpt4TokenEncoder, err := tiktoken.EncodingForModel("gpt-4") gpt4TokenEncoder, err := tiktoken.EncodingForModel("gpt-4")
gpt4oTokenEncoder, err := tiktoken.EncodingForModel("gpt-4o")
if err != nil { if err != nil {
common.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error())) common.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
} }
@ -33,7 +34,11 @@ func InitTokenEncoders() {
if strings.HasPrefix(model, "gpt-3.5") { if strings.HasPrefix(model, "gpt-3.5") {
tokenEncoderMap[model] = gpt35TokenEncoder tokenEncoderMap[model] = gpt35TokenEncoder
} else if strings.HasPrefix(model, "gpt-4") { } else if strings.HasPrefix(model, "gpt-4") {
tokenEncoderMap[model] = gpt4TokenEncoder if strings.HasPrefix(model, "gpt-4o") {
tokenEncoderMap[model] = gpt4oTokenEncoder
} else {
tokenEncoderMap[model] = gpt4TokenEncoder
}
} else { } else {
tokenEncoderMap[model] = nil tokenEncoderMap[model] = nil
} }