From af8c249e3f50fc77d57b5fc797b76d7edf73d0d9 Mon Sep 17 00:00:00 2001
From: MartialBE <me@xiao5.info>
Date: Sun, 26 May 2024 00:44:10 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=96=20chore:=20Allow=20disable=20token?=
 =?UTF-8?q?=20encoders=20(#226)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 common/constants.go |  1 +
 common/token.go     | 12 +++++++++++-
 config.example.yaml |  1 +
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/common/constants.go b/common/constants.go
index 28911a9f..d92ffffd 100644
--- a/common/constants.go
+++ b/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
diff --git a/common/token.go b/common/token.go
index 14bd5ad7..77ba8443 100644
--- a/common/token.go
+++ b/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))
diff --git a/config.example.yaml b/config.example.yaml
index 1f205223..6fd19ed9 100644
--- a/config.example.yaml
+++ b/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