update:num token

This commit is contained in:
Sakurasan
2023-07-08 20:56:54 +08:00
parent c3000c401c
commit c47b540664

View File

@@ -737,38 +737,53 @@ func fetchResponseContent(ctx *gin.Context, responseBody *bufio.Reader) <-chan s
return contentCh return contentCh
} }
func NumTokensFromMessages(messages []openai.ChatCompletionMessage, model string) (num_tokens int) { func NumTokensFromMessages(messages []openai.ChatCompletionMessage, model string) (numTokens int) {
tkm, err := tiktoken.EncodingForModel(model) tkm, err := tiktoken.EncodingForModel(model)
if err != nil { if err != nil {
err = fmt.Errorf("EncodingForModel: %v", err) err = fmt.Errorf("EncodingForModel: %v", err)
fmt.Println(err) log.Println(err)
return return
} }
var tokens_per_message int var tokensPerMessage, tokensPerName int
var tokens_per_name int
if model == "gpt-3.5-turbo-0301" || model == "gpt-3.5-turbo" { switch model {
tokens_per_message = 4 case "gpt-3.5-turbo-0613",
tokens_per_name = -1 "gpt-3.5-turbo-16k-0613",
} else if model == "gpt-4-0314" || model == "gpt-4" { "gpt-4-0314",
tokens_per_message = 3 "gpt-4-32k-0314",
tokens_per_name = 1 "gpt-4-0613",
} else { "gpt-4-32k-0613":
fmt.Println("Warning: model not found. Using cl100k_base encoding.") tokensPerMessage = 3
tokens_per_message = 3 tokensPerName = 1
tokens_per_name = 1 case "gpt-3.5-turbo-0301":
tokensPerMessage = 4 // every message follows <|start|>{role/name}\n{content}<|end|>\n
tokensPerName = -1 // if there's a name, the role is omitted
default:
if strings.Contains(model, "gpt-3.5-turbo") {
log.Println("warning: gpt-3.5-turbo may update over time. Returning num tokens assuming gpt-3.5-turbo-0613.")
return NumTokensFromMessages(messages, "gpt-3.5-turbo-0613")
} else if strings.Contains(model, "gpt-4") {
log.Println("warning: gpt-4 may update over time. Returning num tokens assuming gpt-4-0613.")
return NumTokensFromMessages(messages, "gpt-4-0613")
} else {
err = fmt.Errorf("warning: unknown model [%s]. Use default calculation method converted tokens.", model)
log.Println(err)
return NumTokensFromMessages(messages, "gpt-3.5-turbo-0613")
}
} }
for _, message := range messages { for _, message := range messages {
num_tokens += tokens_per_message numTokens += tokensPerMessage
num_tokens += len(tkm.Encode(message.Content, nil, nil)) numTokens += len(tkm.Encode(message.Content, nil, nil))
// num_tokens += len(tkm.Encode(message.Role, nil, nil)) numTokens += len(tkm.Encode(message.Role, nil, nil))
numTokens += len(tkm.Encode(message.Name, nil, nil))
if message.Name != "" { if message.Name != "" {
num_tokens += tokens_per_name numTokens += tokensPerName
} }
} }
num_tokens += 3 numTokens += 3
return num_tokens return numTokens
} }
func NumTokensFromStr(messages string, model string) (num_tokens int) { func NumTokensFromStr(messages string, model string) (num_tokens int) {