Count system prompt tokens (#850)

* Count sysmte prompt tokens

* show error only once

* fix

* fix css

* simplify

* fix

* simplify

* Revert "simplify"

This reverts commit 61c0a2281f9795c4a4351d9ad473993a3c01c59f.

* `model.tokenizer` config & fix reactivity issues

* rm gated tokenizer

* use `truncate`

---------

Co-authored-by: Nathan Sarrazin <sarrazin.nathan@gmail.com>

.env.template

@@ -7,6 +7,7 @@ MODELS=`[
     "logoUrl": "https://huggingface.co/datasets/huggingchat/models-logo/resolve/main/mistral-logo.png",
     "websiteUrl" : "https://mistral.ai/news/mixtral-of-experts/",
     "modelUrl": "https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "tokenizer": "mistralai/Mixtral-8x7B-Instruct-v0.1",
     "preprompt" : "",
     "chatPromptTemplate": "<s> {{#each messages}}{{#ifUser}}[INST]{{#if @first}}{{#if @root.preprompt}}{{@root.preprompt}}\n{{/if}}{{/if}} {{content}} [/INST]{{/ifUser}}{{#ifAssistant}} {{content}}</s> {{/ifAssistant}}{{/each}}",
     "parameters" : {
@@ -63,7 +64,6 @@ MODELS=`[
       "description": "The latest and biggest model from Meta, fine-tuned for chat.",
       "logoUrl": "https://huggingface.co/datasets/huggingchat/models-logo/resolve/main/meta-logo.png",
       "websiteUrl": "https://ai.meta.com/llama/",
-      "modelUrl": "https://huggingface.co/meta-llama/Llama-2-70b-chat-hf",
       "preprompt": " ",
       "chatPromptTemplate" : "<s>[INST] <<SYS>>\n{{preprompt}}\n<</SYS>>\n\n{{#each messages}}{{#ifUser}}{{content}} [/INST] {{/ifUser}}{{#ifAssistant}}{{content}} </s><s>[INST] {{/ifAssistant}}{{/each}}",
       "promptExamples": [
@@ -94,6 +94,7 @@ MODELS=`[
       "logoUrl": "https://huggingface.co/datasets/huggingchat/models-logo/resolve/main/nous-logo.png",
       "websiteUrl" : "https://nousresearch.com/",
       "modelUrl": "https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
+      "tokenizer": "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
       "chatPromptTemplate" : "{{#if @root.preprompt}}<|im_start|>system\n{{@root.preprompt}}<|im_end|>\n{{/if}}{{#each messages}}{{#ifUser}}<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n{{/ifUser}}{{#ifAssistant}}{{content}}<|im_end|>\n{{/ifAssistant}}{{/each}}",
       "promptExamples": [
         {
@@ -155,6 +156,7 @@ MODELS=`[
       "logoUrl": "https://huggingface.co/datasets/huggingchat/models-logo/resolve/main/mistral-logo.png",
       "websiteUrl": "https://mistral.ai/news/announcing-mistral-7b/",
       "modelUrl": "https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1",
+      "tokenizer": "mistralai/Mistral-7B-Instruct-v0.1",
       "preprompt": "",
       "chatPromptTemplate" : "<s>{{#each messages}}{{#ifUser}}[INST] {{#if @first}}{{#if @root.preprompt}}{{@root.preprompt}}\n{{/if}}{{/if}}{{content}} [/INST]{{/ifUser}}{{#ifAssistant}}{{content}}</s>{{/ifAssistant}}{{/each}}",
       "parameters": {
@@ -187,6 +189,7 @@ MODELS=`[
       "logoUrl": "https://huggingface.co/datasets/huggingchat/models-logo/resolve/main/mistral-logo.png",
       "websiteUrl": "https://mistral.ai/news/announcing-mistral-7b/",
       "modelUrl": "https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2",
+      "tokenizer": "mistralai/Mistral-7B-Instruct-v0.2",
       "preprompt": "",
       "chatPromptTemplate" : "<s>{{#each messages}}{{#ifUser}}[INST] {{#if @first}}{{#if @root.preprompt}}{{@root.preprompt}}\n{{/if}}{{/if}}{{content}} [/INST]{{/ifUser}}{{#ifAssistant}}{{content}}</s>{{/ifAssistant}}{{/each}}",
       "parameters": {
@@ -218,6 +221,7 @@ MODELS=`[
     "logoUrl": "https://huggingface.co/datasets/huggingchat/models-logo/resolve/main/openchat-logo.png",
     "websiteUrl": "https://huggingface.co/openchat/openchat-3.5-0106",
     "modelUrl": "https://huggingface.co/openchat/openchat-3.5-0106",
+    "tokenizer": "openchat/openchat-3.5-0106",
     "preprompt": "",
     "chatPromptTemplate" : "<s>{{#each messages}}{{#ifUser}}GPT4 Correct User: {{#if @first}}{{#if @root.preprompt}}{{@root.preprompt}}\n{{/if}}{{/if}}{{content}}<|end_of_turn|>GPT4 Correct Assistant:{{/ifUser}}{{#ifAssistant}}{{content}}<|end_of_turn|>{{/ifAssistant}}{{/each}}",
     "parameters": {

src/lib/components/AssistantSettings.svelte

@@ -12,6 +12,7 @@
 
 	import { useSettingsStore } from "$lib/stores/settings";
 	import { isHuggingChat } from "$lib/utils/isHuggingChat";
+	import TokensCounter from "./TokensCounter.svelte";
 
 	type ActionData = {
 		error: boolean;
@@ -28,8 +29,10 @@
 	export let models: Model[] = [];
 
 	let files: FileList | null = null;
-
 	const settings = useSettingsStore();
+	let modelId =
+		assistant?.modelId ?? models.find((_model) => _model.id === $settings.activeModel)?.name;
+	let systemPrompt = assistant?.preprompt ?? "";
 
 	let compress: typeof readAndCompressImage | null = null;
 
@@ -238,7 +241,11 @@
 
 			<label>
 				<div class="mb-1 font-semibold">Model</div>
-				<select name="modelId" class="w-full rounded-lg border-2 border-gray-200 bg-gray-100 p-2">
+				<select
+					name="modelId"
+					class="w-full rounded-lg border-2 border-gray-200 bg-gray-100 p-2"
+					bind:value={modelId}
+				>
 					{#each models.filter((model) => !model.unlisted) as model}
 						<option
 							value={model.id}
@@ -390,12 +397,25 @@
 
 		<div class="col-span-1 flex h-full flex-col">
 			<span class="mb-1 text-sm font-semibold"> Instructions (system prompt) </span>
-			<textarea
-				name="preprompt"
-				class="mb-20 min-h-[8lh] flex-1 rounded-lg border-2 border-gray-200 bg-gray-100 p-2 text-sm"
-				placeholder="You'll act as..."
-				value={assistant?.preprompt ?? ""}
-			/>
+			<div class="relative mb-20 flex min-h-[8lh] flex-1 grow flex-col">
+				<textarea
+					name="preprompt"
+					class="flex-1 rounded-lg border-2 border-gray-200 bg-gray-100 p-2 text-sm"
+					placeholder="You'll act as..."
+					bind:value={systemPrompt}
+				/>
+				{#if modelId}
+					{@const model = models.find((_model) => _model.id === modelId)}
+					{#if model?.tokenizer && systemPrompt}
+						<TokensCounter
+							classNames="absolute bottom-2 right-2"
+							prompt={systemPrompt}
+							modelTokenizer={model.tokenizer}
+							truncate={model?.parameters?.truncate}
+						/>
+					{/if}
+				{/if}
+			</div>
 			<p class="text-xs text-red-500">{getError("preprompt", form)}</p>
 		</div>
 	</div>

src/lib/components/TokensCounter.svelte

@@ -0,0 +1,48 @@
+<script lang="ts">
+	import type { Model } from "$lib/types/Model";
+	import { AutoTokenizer, PreTrainedTokenizer } from "@xenova/transformers";
+
+	export let classNames = "";
+	export let prompt = "";
+	export let modelTokenizer: Exclude<Model["tokenizer"], undefined>;
+	export let truncate: number | undefined = undefined;
+
+	let tokenizer: PreTrainedTokenizer | undefined = undefined;
+
+	async function getTokenizer(_modelTokenizer: Exclude<Model["tokenizer"], undefined>) {
+		if (typeof _modelTokenizer === "string") {
+			// return auto tokenizer
+			return await AutoTokenizer.from_pretrained(_modelTokenizer);
+		}
+		{
+			// construct & return pretrained tokenizer
+			const { tokenizerUrl, tokenizerConfigUrl } = _modelTokenizer satisfies {
+				tokenizerUrl: string;
+				tokenizerConfigUrl: string;
+			};
+			const tokenizerJSON = await (await fetch(tokenizerUrl)).json();
+			const tokenizerConfig = await (await fetch(tokenizerConfigUrl)).json();
+			return new PreTrainedTokenizer(tokenizerJSON, tokenizerConfig);
+		}
+	}
+
+	async function tokenizeText(_prompt: string) {
+		if (!tokenizer) {
+			return;
+		}
+		const { input_ids } = await tokenizer(_prompt);
+		return input_ids.size;
+	}
+
+	$: (async () => {
+		tokenizer = await getTokenizer(modelTokenizer);
+	})();
+</script>
+
+{#if tokenizer}
+	{#await tokenizeText(prompt) then nTokens}
+		<p class="text-sm opacity-60 hover:opacity-80 {classNames}">
+			{nTokens}{truncate ? `/${truncate}` : ""}
+		</p>
+	{/await}
+{/if}

src/lib/server/models.ts

@@ -28,6 +28,15 @@ const modelConfig = z.object({
 	logoUrl: z.string().url().optional(),
 	websiteUrl: z.string().url().optional(),
 	modelUrl: z.string().url().optional(),
+	tokenizer: z
+		.union([
+			z.string(),
+			z.object({
+				tokenizerUrl: z.string().url(),
+				tokenizerConfigUrl: z.string().url(),
+			}),
+		])
+		.optional(),
 	datasetName: z.string().min(1).optional(),
 	datasetUrl: z.string().url().optional(),
 	userMessageToken: z.string().default(""),

src/lib/types/Model.ts

@@ -12,6 +12,7 @@ export type Model = Pick<
 	| "description"
 	| "logoUrl"
 	| "modelUrl"
+	| "tokenizer"
 	| "datasetUrl"
 	| "preprompt"
 	| "multimodal"

src/routes/+layout.server.ts

@@ -158,6 +158,7 @@ export const load: LayoutServerLoad = async ({ locals, depends }) => {
 			name: model.name,
 			websiteUrl: model.websiteUrl,
 			modelUrl: model.modelUrl,
+			tokenizer: model.tokenizer,
 			datasetName: model.datasetName,
 			datasetUrl: model.datasetUrl,
 			displayName: model.displayName,

src/routes/api/models/+server.ts

@@ -6,6 +6,7 @@ export async function GET() {
 		name: model.name,
 		websiteUrl: model.websiteUrl,
 		modelUrl: model.modelUrl,
+		tokenizer: model.tokenizer,
 		datasetName: model.datasetName,
 		datasetUrl: model.datasetUrl,
 		displayName: model.displayName,

src/routes/settings/(nav)/[...model]/+page.svelte

@@ -5,6 +5,7 @@
 	import type { BackendModel } from "$lib/server/models";
 	import { useSettingsStore } from "$lib/stores/settings";
 	import CopyToClipBoardBtn from "$lib/components/CopyToClipBoardBtn.svelte";
+	import TokensCounter from "$lib/components/TokensCounter.svelte";
 	import CarbonArrowUpRight from "~icons/carbon/arrow-up-right";
 	import CarbonLink from "~icons/carbon/link";
 
@@ -99,7 +100,7 @@
 		{isActive ? "Active model" : "Activate"}
 	</button>
 
-	<div class="flex w-full flex-col gap-2">
+	<div class="relative flex w-full flex-col gap-2">
 		<div class="flex w-full flex-row content-between">
 			<h3 class="mb-1.5 text-lg font-semibold text-gray-800">System Prompt</h3>
 			{#if hasCustomPreprompt}
@@ -117,5 +118,13 @@
 			class="w-full resize-none rounded-md border-2 bg-gray-100 p-2"
 			bind:value={$settings.customPrompts[$page.params.model]}
 		/>
+		{#if model.tokenizer && $settings.customPrompts[$page.params.model]}
+			<TokensCounter
+				classNames="absolute bottom-2 right-2"
+				prompt={$settings.customPrompts[$page.params.model]}
+				modelTokenizer={model.tokenizer}
+				truncate={model?.parameters?.truncate}
+			/>
+		{/if}
 	</div>
 </div>