Count system prompt tokens (#850)

* Count sysmte prompt tokens * show error only once * fix * fix css * simplify * fix * simplify * Revert "simplify" This reverts commit 61c0a22. * `model.tokenizer` config & fix reactivity issues * rm gated tokenizer * use `truncate` --------- Co-authored-by: Nathan Sarrazin <[email protected]>
huggingface · Mar 22, 2024 · 50d8483 · 50d8483
1 parent c6129c3
commit 50d8483
Show file tree

Hide file tree

Showing 8 changed files with 103 additions and 10 deletions.
diff --git a/.env.template b/.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": {

diff --git a/src/lib/components/AssistantSettings.svelte b/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>

diff --git a/src/lib/components/TokensCounter.svelte b/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}
diff --git a/src/lib/server/models.ts b/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(""),

diff --git a/src/lib/types/Model.ts b/src/lib/types/Model.ts
@@ -12,6 +12,7 @@ export type Model = Pick<
  | "description"
  | "logoUrl"
  | "modelUrl"
+ | "tokenizer"
  | "datasetUrl"
  | "preprompt"
  | "multimodal"

diff --git a/src/routes/+layout.server.ts b/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,

diff --git a/src/routes/api/models/+server.ts b/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,

diff --git a/src/routes/settings/(nav)/[...model]/+page.svelte b/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>