⚡️ Limit the number of tokens sent to the backend (#93)

.env

@@ -6,6 +6,7 @@ MONGODB_DB_NAME=chat-ui
 HF_TOKEN=#your huggingface token here
 COOKIE_NAME=hf-chat
 
+PUBLIC_MAX_INPUT_TOKENS=1024
 PUBLIC_ORIGIN=#https://hf.co
 PUBLIC_MODEL_ENDPOINT=https://api-inference.huggingface.co/models/OpenAssistant/oasst-sft-6-llama-30b
 PUBLIC_MODEL_NAME=OpenAssistant/oasst-sft-6-llama-30b # public facing link

src/lib/buildPrompt.ts

@@ -1,5 +1,6 @@
 import {
 	PUBLIC_ASSISTANT_MESSAGE_TOKEN,
+	PUBLIC_MAX_INPUT_TOKENS,
 	PUBLIC_SEP_TOKEN,
 	PUBLIC_USER_MESSAGE_TOKEN,
 } from "$env/static/public";
@@ -11,7 +12,7 @@ import type { Message } from "./types/Message";
  * <|assistant|>hi<|endoftext|><|prompter|>hello<|endoftext|><|assistant|>
  */
 export function buildPrompt(messages: Message[]): string {
-	return (
+	const prompt =
 		messages
 			.map(
 				(m) =>
@@ -20,6 +21,8 @@ export function buildPrompt(messages: Message[]): string {
 						: PUBLIC_ASSISTANT_MESSAGE_TOKEN + m.content) +
 					(m.content.endsWith(PUBLIC_SEP_TOKEN) ? "" : PUBLIC_SEP_TOKEN)
 			)
-			.join("") + PUBLIC_ASSISTANT_MESSAGE_TOKEN
-	);
+			.join("") + PUBLIC_ASSISTANT_MESSAGE_TOKEN;
+
+	// Not super precise, but it's truncated in the model's backend anyway
+	return prompt.split(" ").slice(-parseInt(PUBLIC_MAX_INPUT_TOKENS)).join(" ");
 }

src/routes/conversation/[id]/+page.svelte

@@ -8,7 +8,7 @@
 	import { invalidate } from "$app/navigation";
 	import { base } from "$app/paths";
 	import { trimSuffix } from "$lib/utils/trimSuffix";
-	import { PUBLIC_SEP_TOKEN } from "$env/static/public";
+	import { PUBLIC_SEP_TOKEN, PUBLIC_MAX_INPUT_TOKENS } from "$env/static/public";
 	import { trimPrefix } from "$lib/utils/trimPrefix";
 	import { shareConversation } from "$lib/shareConversation";
 	import { UrlDependency } from "$lib/types/UrlDependency";
@@ -41,7 +41,7 @@
 					repetition_penalty: 1.2,
 					top_k: 50,
 					// @ts-ignore
-					truncate: 1024,
+					truncate: parseInt(PUBLIC_MAX_INPUT_TOKENS),
 					watermark: false,
 					max_new_tokens: 1024,
 					stop: ["<|endoftext|>"],

src/routes/conversation/[id]/summarize/+server.ts

@@ -1,5 +1,5 @@
 import { HF_TOKEN } from "$env/static/private";
-import { PUBLIC_MODEL_ENDPOINT } from "$env/static/public";
+import { PUBLIC_MAX_INPUT_TOKENS, PUBLIC_MODEL_ENDPOINT } from "$env/static/public";
 import { buildPrompt } from "$lib/buildPrompt";
 import { collections } from "$lib/server/database.js";
 import { textGeneration } from "@huggingface/inference";
@@ -33,6 +33,7 @@ export async function POST({ params, locals, fetch }) {
 		top_k: 50,
 		watermark: false,
 		max_new_tokens: 1024,
+		truncate: parseInt(PUBLIC_MAX_INPUT_TOKENS),
 		stop: ["<|endoftext|>"],
 		return_full_text: false,
 	};