From e9d3b4654f919c84457231a03061964e9d72ea5e Mon Sep 17 00:00:00 2001 From: MartialBE Date: Tue, 14 May 2024 21:55:24 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=96=20chore:=20add=20gpt-4o=20tiktoken?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- common/token.go | 10 +++++++++- go.mod | 2 +- go.sum | 4 ++-- 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/common/token.go b/common/token.go index 40ef8733..14bd5ad7 100644 --- a/common/token.go +++ b/common/token.go @@ -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 { diff --git a/go.mod b/go.mod index db1ed215..6e4afad6 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 1647ae52..85532af6 100644 --- a/go.sum +++ b/go.sum @@ -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=