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

🔖 chore: add gpt-4o tiktoken

Browse Source
This commit is contained in:
MartialBE 2024-05-14 21:55:24 +08:00
parent d48868586c
commit e9d3b4654f
No known key found for this signature in database
GPG Key ID: 27C0267EC84B0A5C
3 changed files with 12 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
common/token.go
View File

@ -9,12 +9,13 @@ import (
"one-api/common/image"
"one-api/types"
"github.com/pkoukk/tiktoken-go"
"github.com/MartialBE/tiktoken-go"
)
var tokenEncoderMap = map[string]*tiktoken.Tiktoken{}
var gpt35TokenEncoder *tiktoken.Tiktoken
var gpt4TokenEncoder *tiktoken.Tiktoken
var gpt4oTokenEncoder *tiktoken.Tiktoken
func InitTokenEncoders() {
SysLog("initializing token encoders")
@ -29,6 +30,11 @@ func InitTokenEncoders() {
FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
}
gpt4oTokenEncoder, err = tiktoken.EncodingForModel("gpt-4o")
if err != nil {
FatalLog(fmt.Sprintf("failed to get gpt-4o token encoder: %s", err.Error()))
}
SysLog("token encoders initialized")
}
@ -40,6 +46,8 @@ func getTokenEncoder(model string) *tiktoken.Tiktoken {
if strings.HasPrefix(model, "gpt-3.5") {
tokenEncoder = gpt35TokenEncoder
} else if strings.HasPrefix(model, "gpt-4o") {
tokenEncoder = gpt4oTokenEncoder
} else if strings.HasPrefix(model, "gpt-4") {
tokenEncoder = gpt4TokenEncoder
} else {

2
go.mod
View File

@ -4,6 +4,7 @@ module one-api
go 1.18
require (
github.com/MartialBE/tiktoken-go v0.1.7
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.1
github.com/aws/smithy-go v1.20.1
github.com/gin-contrib/cors v1.7.0
@ -18,7 +19,6 @@ require (
github.com/gomarkdown/markdown v0.0.0-20240328165702-4d01890c35c0
github.com/google/uuid v1.6.0
github.com/gorilla/websocket v1.5.1
github.com/pkoukk/tiktoken-go v0.1.6
github.com/shopspring/decimal v1.4.0
github.com/spf13/viper v1.18.2
github.com/stretchr/testify v1.9.0

4
go.sum
View File

@ -1,5 +1,7 @@
filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
github.com/MartialBE/tiktoken-go v0.1.7 h1:QwYqiONL1+kc/bYtNIbfQ//bcRqyw8ETZ2HKKZTQv8A=
github.com/MartialBE/tiktoken-go v0.1.7/go.mod h1:BTUpeOlungENRv13+nWc4IwWbBJvvR/8gOUTiXsrEXg=
github.com/PaulSonOfLars/gotgbot/v2 v2.0.0-rc.25 h1:VCZg3OsKY19PcXBRRYk2ExeZ3mC8Hm4LqcXcINuFyY4=
github.com/PaulSonOfLars/gotgbot/v2 v2.0.0-rc.25/go.mod h1:kL1v4iIjlalwm3gCYGvF4NLa3hs+aKEfRkNJvj4aoDU=
github.com/antonlindstrom/pgstore v0.0.0-20200229204646-b08ebf1105e0/go.mod h1:2Ti6VUHVxpC0VSmTZzEvpzysnaGAfGBOoMIz5ykPyyw=
@ -166,8 +168,6 @@ github.com/pelletier/go-toml/v2 v2.1.1 h1:LWAJwfNvjQZCFIDKWYQaM62NcYeYViCmWIwmOS
github.com/pelletier/go-toml/v2 v2.1.1/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
github.com/pkoukk/tiktoken-go v0.1.6 h1:JF0TlJzhTbrI30wCvFuiw6FzP2+/bR+FIxUdgEAcUsw=
github.com/pkoukk/tiktoken-go v0.1.6/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=