修复gpt-3.5-turbo-0301 无法计算token的问题

3 years ago · 96334cc572
parent 9fd3ac3923
commit 96334cc572
1 changed files with 6 additions and 0 deletions
--- a/main.go
+++ b/main.go
@ -150,6 +150,9 @@ func InitRedis() *redis.Client {

 // 计算Messages中的token数量
 func numTokensFromMessages(messages []Message, model string) int {
+	if model == "gpt-3.5-turbo-0301" {
+		model = "gpt-3.5-turbo"
+	}
 	tkm, err := tiktoken.EncodingForModel(model)
 	if err != nil {
 		err = fmt.Errorf("getEncoding: %v", err)
@ -166,6 +169,9 @@ func numTokensFromMessages(messages []Message, model string) int {

 // 计算String中的token数量
 func numTokensFromString(msg string, model string) int {
+	if model == "gpt-3.5-turbo-0301" {
+		model = "gpt-3.5-turbo"
+	}
 	tkm, err := tiktoken.EncodingForModel(model)
 	if err != nil {
 		err = fmt.Errorf("getEncoding: %v", err)