feat: new token count function

This commit is contained in:
Yidadaa 2023-06-15 00:14:38 +08:00
parent 8590750e4c
commit 76fdd047e7
2 changed files with 24 additions and 1 deletions

View File

@ -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
app/utils/token.ts Normal file
View File

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