Browse Source

feat: new token count function

Yidadaa 1 year ago
parent
commit
76fdd047e7
2 changed files with 24 additions and 1 deletions
  1. 2 1
      app/store/chat.ts
  2. 22 0
      app/utils/token.ts

+ 2 - 1
app/store/chat.ts

@@ -11,6 +11,7 @@ import { StoreKey } from "../constant";
 import { api, RequestMessage } from "../client/api";
 import { ChatControllerPool } from "../client/controller";
 import { prettyObject } from "../utils/format";
+import { estimateTokenLength } from "../utils/token";
 
 export type ChatMessage = RequestMessage & {
   date: string;
@@ -102,7 +103,7 @@ interface ChatStore {
 }
 
 function countMessages(msgs: ChatMessage[]) {
-  return msgs.reduce((pre, cur) => pre + cur.content.length, 0);
+  return msgs.reduce((pre, cur) => pre + estimateTokenLength(cur.content), 0);
 }
 
 export const useChatStore = create<ChatStore>()(

+ 22 - 0
app/utils/token.ts

@@ -0,0 +1,22 @@
+export function estimateTokenLength(input: string): number {
+  let tokenLength = 0;
+
+  for (let i = 0; i < input.length; i++) {
+    const charCode = input.charCodeAt(i);
+
+    if (charCode < 128) {
+      // ASCII character
+      if (charCode <= 122 && charCode >= 65) {
+        // a-Z
+        tokenLength += 0.25;
+      } else {
+        tokenLength += 0.5;
+      }
+    } else {
+      // Unicode character
+      tokenLength += 1.5;
+    }
+  }
+
+  return tokenLength;
+}