Browse Source

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

JustSong 2 years ago
parent
commit
fdb2cccf65
2 changed files with 19 additions and 0 deletions
  1. 18 0
      controller/relay-utils.go
  2. 1 0
      main.go

+ 18 - 0
controller/relay-utils.go

@@ -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

+ 1 - 0
main.go

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