Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: update countTokenMessages (#238)

Browse Source
This commit is contained in:
quzard 2023-07-03 09:42:34 +08:00 committed by GitHub
parent c90c0ecef4
commit b09daf5ec1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
controller/relay-utils.go
View File

@ -4,7 +4,6 @@ import (
"fmt"
"github.com/pkoukk/tiktoken-go"
"one-api/common"
"strings"
)
var tokenEncoderMap = map[string]*tiktoken.Tiktoken{}
@ -34,12 +33,9 @@ func countTokenMessages(messages []Message, model string) int {
// Every message follows <|start|>{role/name}\n{content}<|end|>\n
var tokensPerMessage int
var tokensPerName int
if strings.HasPrefix(model, "gpt-3.5") {
if model == "gpt-3.5-turbo-0301" {
tokensPerMessage = 4
tokensPerName = -1 // If there's a name, the role is omitted
} else if strings.HasPrefix(model, "gpt-4") {
tokensPerMessage = 3
tokensPerName = 1
} else {
tokensPerMessage = 3
tokensPerName = 1