🔖 chore: Allow disable token encoders (#226)

common/constants.go

@@ -90,6 +90,7 @@ var AutomaticEnableChannelEnabled = false
 var QuotaRemindThreshold = 1000
 var PreConsumedQuota = 500
 var ApproximateTokenEnabled = false
+var DISABLE_TOKEN_ENCODERS = false
 var RetryTimes = 0
 var DefaultChannelWeight = uint(1)
 var RetryCooldownSeconds = 5

common/token.go

@@ -10,6 +10,7 @@ import (
 	"one-api/types"
 
 	"github.com/MartialBE/tiktoken-go"
+	"github.com/spf13/viper"
 )
 
 var tokenEncoderMap = map[string]*tiktoken.Tiktoken{}
@@ -18,6 +19,11 @@ var gpt4TokenEncoder *tiktoken.Tiktoken
 var gpt4oTokenEncoder *tiktoken.Tiktoken
 
 func InitTokenEncoders() {
+	if viper.GetBool("disable_token_encoders") {
+		DISABLE_TOKEN_ENCODERS = true
+		SysLog("token encoders disabled")
+		return
+	}
 	SysLog("initializing token encoders")
 	var err error
 	gpt35TokenEncoder, err = tiktoken.EncodingForModel("gpt-3.5-turbo")
@@ -39,6 +45,10 @@ func InitTokenEncoders() {
 }
 
 func getTokenEncoder(model string) *tiktoken.Tiktoken {
+	if DISABLE_TOKEN_ENCODERS {
+		return nil
+	}
+
 	tokenEncoder, ok := tokenEncoderMap[model]
 	if ok {
 		return tokenEncoder
@@ -64,7 +74,7 @@ func getTokenEncoder(model string) *tiktoken.Tiktoken {
 }
 
 func getTokenNum(tokenEncoder *tiktoken.Tiktoken, text string) int {
-	if ApproximateTokenEnabled {
+	if DISABLE_TOKEN_ENCODERS || ApproximateTokenEnabled {
 		return int(float64(len(text)) * 0.38)
 	}
 	return len(tokenEncoder.Encode(text, nil, nil))