perf: initialize all token encoder when starting (close #459, close $460)

This commit is contained in:
JustSong 2023-08-26 13:02:02 +08:00
parent a3e267df7e
commit fdb2cccf65
2 changed files with 19 additions and 0 deletions

View File

@ -15,6 +15,24 @@ var stopFinishReason = "stop"
var tokenEncoderMap = map[string]*tiktoken.Tiktoken{}
func InitTokenEncoders() {
common.SysLog("initializing token encoders")
fallbackTokenEncoder, err := tiktoken.EncodingForModel("gpt-3.5-turbo")
if err != nil {
common.FatalLog(fmt.Sprintf("failed to get fallback token encoder: %s", err.Error()))
}
for model, _ := range common.ModelRatio {
tokenEncoder, err := tiktoken.EncodingForModel(model)
if err != nil {
common.SysError(fmt.Sprintf("using fallback encoder for model %s", model))
tokenEncoderMap[model] = fallbackTokenEncoder
continue
}
tokenEncoderMap[model] = tokenEncoder
}
common.SysLog("token encoders initialized")
}
func getTokenEncoder(model string) *tiktoken.Tiktoken {
if tokenEncoder, ok := tokenEncoderMap[model]; ok {
return tokenEncoder

View File

@ -77,6 +77,7 @@ func main() {
}
go controller.AutomaticallyTestChannels(frequency)
}
controller.InitTokenEncoders()
// Initialize HTTP server
server := gin.Default()