🔖 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))

config.example.yaml

@@ -4,6 +4,7 @@ gin_mode: "release" # gin 模式，可选值为 "release" 或 "debug"，默认
 debug: false # 是否启用调试模式，启用后将输出更多日志信息。
 log_dir: "./logs" # 日志目录
 session_secret: "" # 会话密钥，未设置则使用随机值。
+disable_token_encoders: false # 是否禁用 token 编码器计算tokens。启用后 内存占用可减少 40MB 左右，但是stream模式下tokens计算不准确
 
 # 数据库设置
 sql_dsn: "" # 设置之后将使用指定数据库而非 SQLite，请使用 MySQL 或 PostgreSQL