feat: new token count function

jimye · Jun 14, 2023 · 76fdd04 · 76fdd04
1 parent 8590750
commit 76fdd04
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/app/store/chat.ts b/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>()(

diff --git a/app/utils/token.ts b/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;
+}