Trim system prompt for delta messages

src/chat.ts

@@ -43,7 +43,7 @@ async function main(dir: string) {
     const question = await rl.question('You> ')
     messages.push({role: 'user', content: question});
     process.stdout.write('Assistant> ');
-    const reply = await talk(tokenizer, model, kvCache, messages.slice(-1));
+    const reply = await talk(tokenizer, model, kvCache, messages.slice(-1), messages.length == 1);
     messages.push({role: 'assistant', content: reply});
   }
 }
@@ -52,9 +52,12 @@ async function main(dir: string) {
 async function talk(tokenizer: Tokenizer,
                     model: BaseModel,
                     kvCache: BaseKVCache[],
-                    messages: Message[]) {
+                    messages: Message[],
+                    firstMessage: boolean) {
   // Translate the messages to tokens.
-  const promptTokens = tokenizer.applyChatTemplate(messages);
+  // Note that some chat templates add a system prompt automatically and we need
+  // to trim it when generating tokens for only new messages.
+  const promptTokens = tokenizer.applyChatTemplate(messages, {trimSystemPrompt: !firstMessage});
 
   // Predict next tokens.
   let tokens: number[] = [];

src/tokenizer.ts

@@ -9,13 +9,21 @@ export interface Message {
   content: string;
 }
 
+/**
+ * Options for chat template.
+ */
+export interface ChatTemplateOptions {
+  trimSystemPrompt?: boolean;
+}
+
 /**
  * Wraps the tokenizer of transformers.js.
  */
 export class Tokenizer {
   bosToken: number;
   eosToken: number;
   private tokenizer: ReturnType<typeof TokenizerLoader.fromPreTrained>;
+  private systemPromptLength?: number;
 
   constructor(dir: string) {
     this.tokenizer = TokenizerLoader.fromPreTrained({
@@ -41,11 +49,25 @@ export class Tokenizer {
     return this.tokenizer.decode(tokens);
   }
 
-  applyChatTemplate(messages: Message[]): number[] {
-    return this.tokenizer.apply_chat_template(messages, {
+  applyChatTemplate(messages: Message[],
+                    {
+                      trimSystemPrompt = false,
+                    }: ChatTemplateOptions = {}): number[] {
+    if (trimSystemPrompt && this.systemPromptLength === undefined) {
+      // Get the automatically inserted system prompt by passing empty messages.
+      const systemPrompt = this.tokenizer.apply_chat_template([], {
+        add_generation_prompt: false,
+        tools: null,
+      } as unknown) as number[];
+      this.systemPromptLength = systemPrompt.length;
+    }
+    const tokens = this.tokenizer.apply_chat_template(messages, {
       add_generation_prompt: true,
       // https://github.com/xenova/transformers.js/issues/879
       tools: null,
     } as unknown) as number[];
+    if (trimSystemPrompt)
+      return tokens.slice(this.systemPromptLength);
+    return tokens;
   }
 }