xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-13 03:43:44 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor: refactor relay part

Browse Source
This commit is contained in:
JustSong
2024-01-21 22:56:20 +08:00
parent e2ed0399f0
commit b373882814
64 changed files with 1092 additions and 793 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
relay/channel/openai/token.go
View File

@@ -7,6 +7,7 @@ import (
"math"
"one-api/common"
"one-api/common/image"
"one-api/common/logger"
"strings"
)
@@ -15,15 +16,15 @@ var tokenEncoderMap = map[string]*tiktoken.Tiktoken{}
var defaultTokenEncoder *tiktoken.Tiktoken
func InitTokenEncoders() {
common.SysLog("initializing token encoders")
logger.SysLog("initializing token encoders")
gpt35TokenEncoder, err := tiktoken.EncodingForModel("gpt-3.5-turbo")
if err != nil {
common.FatalLog(fmt.Sprintf("failed to get gpt-3.5-turbo token encoder: %s", err.Error()))
logger.FatalLog(fmt.Sprintf("failed to get gpt-3.5-turbo token encoder: %s", err.Error()))
}
defaultTokenEncoder = gpt35TokenEncoder
gpt4TokenEncoder, err := tiktoken.EncodingForModel("gpt-4")
if err != nil {
common.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
logger.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
}
for model, _ := range common.ModelRatio {
if strings.HasPrefix(model, "gpt-3.5") {
@@ -34,7 +35,7 @@ func InitTokenEncoders() {
tokenEncoderMap[model] = nil
}
}
common.SysLog("token encoders initialized")
logger.SysLog("token encoders initialized")
}
func getTokenEncoder(model string) *tiktoken.Tiktoken {
@@ -45,7 +46,7 @@ func getTokenEncoder(model string) *tiktoken.Tiktoken {
if ok {
tokenEncoder, err := tiktoken.EncodingForModel(model)
if err != nil {
common.SysError(fmt.Sprintf("failed to get token encoder for model %s: %s, using encoder for gpt-3.5-turbo", model, err.Error()))
logger.SysError(fmt.Sprintf("failed to get token encoder for model %s: %s, using encoder for gpt-3.5-turbo", model, err.Error()))
tokenEncoder = defaultTokenEncoder
}
tokenEncoderMap[model] = tokenEncoder
@@ -99,7 +100,7 @@ func CountTokenMessages(messages []Message, model string) int {
}
imageTokens, err := countImageTokens(url, detail)
if err != nil {
common.SysError("error counting image tokens: " + err.Error())
logger.SysError("error counting image tokens: " + err.Error())
} else {
tokenNum += imageTokens
}