Explorar o código

chore: update InitTokenEncoders (#255)

CaIon hai 1 ano
pai
achega
a3b3e6cc38
Modificáronse 1 ficheiros con 6 adicións e 1 borrados
  1. 6 1
      service/token_counter.go

+ 6 - 1
service/token_counter.go

@@ -26,6 +26,7 @@ func InitTokenEncoders() {
 	}
 	}
 	defaultTokenEncoder = gpt35TokenEncoder
 	defaultTokenEncoder = gpt35TokenEncoder
 	gpt4TokenEncoder, err := tiktoken.EncodingForModel("gpt-4")
 	gpt4TokenEncoder, err := tiktoken.EncodingForModel("gpt-4")
+	gpt4oTokenEncoder, err := tiktoken.EncodingForModel("gpt-4o")
 	if err != nil {
 	if err != nil {
 		common.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
 		common.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
 	}
 	}
@@ -33,7 +34,11 @@ func InitTokenEncoders() {
 		if strings.HasPrefix(model, "gpt-3.5") {
 		if strings.HasPrefix(model, "gpt-3.5") {
 			tokenEncoderMap[model] = gpt35TokenEncoder
 			tokenEncoderMap[model] = gpt35TokenEncoder
 		} else if strings.HasPrefix(model, "gpt-4") {
 		} else if strings.HasPrefix(model, "gpt-4") {
-			tokenEncoderMap[model] = gpt4TokenEncoder
+			if strings.HasPrefix(model, "gpt-4o") {
+				tokenEncoderMap[model] = gpt4oTokenEncoder
+			} else {
+				tokenEncoderMap[model] = gpt4TokenEncoder
+			}
 		} else {
 		} else {
 			tokenEncoderMap[model] = nil
 			tokenEncoderMap[model] = nil
 		}
 		}