diff options
Diffstat (limited to 'packages/shared')
| -rw-r--r-- | packages/shared/package.json | 1 | ||||
| -rw-r--r-- | packages/shared/prompts.ts | 16 |
2 files changed, 9 insertions, 8 deletions
diff --git a/packages/shared/package.json b/packages/shared/package.json index a262267a..478fb77d 100644 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -6,6 +6,7 @@ "type": "module", "dependencies": { "glob": "^11.0.0", + "js-tiktoken": "^1.0.20", "liteque": "^0.3.2", "meilisearch": "^0.37.0", "ollama": "^0.5.14", diff --git a/packages/shared/prompts.ts b/packages/shared/prompts.ts index 40e706b2..b7957dd4 100644 --- a/packages/shared/prompts.ts +++ b/packages/shared/prompts.ts @@ -1,15 +1,15 @@ -// TODO: Use a proper tokenizer +import { getEncoding } from "js-tiktoken"; + +const encoding = getEncoding("o200k_base"); + function calculateNumTokens(text: string) { - return text.split(" ").length; + return encoding.encode(text).length; } function truncateContent(content: string, length: number) { - let words = content.split(" "); - if (words.length > length) { - words = words.slice(0, length); - content = words.join(" "); - } - return content; + const tokens = encoding.encode(content); + const truncatedTokens = tokens.slice(0, length); + return encoding.decode(truncatedTokens); } export function buildImagePrompt(lang: string, customPrompts: string[]) { |
