diff --git a/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts b/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts index 7ad8d990..2c1ee111 100644 --- a/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts +++ b/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts @@ -1,10 +1,13 @@ import {Template} from "@huggingface/jinja"; import {splitText} from "lifecycle-utils"; import { - ChatHistoryItem, ChatModelFunctions, ChatUserMessage, ChatWrapperGenerateContextStateOptions, ChatWrapperGeneratedContextState, + allSegmentTypes, ChatHistoryItem, ChatModelFunctions, ChatModelSegmentType, ChatUserMessage, + ChatWrapperGenerateContextStateOptions, ChatWrapperGeneratedContextState, ChatWrapperGeneratedPrefixTriggersContextState, ChatWrapperSettings, Tokenizer } from "../../types.js"; import {SpecialToken, LlamaText, SpecialTokensText} from "../../utils/LlamaText.js"; +import {getChatWrapperSegmentDefinition} from "../../utils/getChatWrapperSegmentDefinition.js"; +import {includesText} from "../../utils/includesText.js"; import {ChatWrapper} from "../../ChatWrapper.js"; import { fromChatHistoryToIntermediateOpenAiMessages, fromIntermediateToCompleteOpenAiMessages, IntermediateOpenAiMessage, @@ -19,7 +22,8 @@ import { } from "./utils/templateSegmentOptionsToChatWrapperSettings.js"; import {UniqueIdGenerator} from "./utils/UniqueIdGenerator.js"; import { - detectNeedToWrapFunctionArgumentsWithMap, extractFunctionCallSettingsFromJinjaTemplate, ExtractFunctionCallSettingsRenderTemplate + detectNeedToWrapFunctionArgumentsWithMap, extractFunctionCallSettingsFromJinjaTemplate, ExtractFunctionCallSettingsRenderTemplate, + findCommonStartLength } from "./utils/extractFunctionCallSettingsFromJinjaTemplate.js"; import {squashChatHistoryItems} from "./utils/squashChatHistoryItems.js"; import {extractSegmentSettingsFromTokenizerAndChatTemplate} from "./utils/extractSegmentSettingsFromTokenizerAndChatTemplate.js"; @@ -306,13 +310,8 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { const lastJinjaItem = messages.at(-1); let eraseRenderedJinjaFromId: string | undefined; - if (this._endJinjaMessagesWithUserMessage && lastJinjaItem?.role === this.modelRoleName && - typeof lastJinjaItem.content === "string" && - lastJinjaItem.content.length > 0 && - ( - (lastJinjaItem as OpenAiChatAssistantMessage)["tool_calls"] == null || - (lastJinjaItem as OpenAiChatAssistantMessage)["tool_calls"]?.length === 0 - ) + if (this._endJinjaMessagesWithUserMessage && isPlainModelMessage(lastJinjaItem, this.modelRoleName) && + lastJinjaItem.content.length > 0 ) { eraseRenderedJinjaFromId = lastJinjaItem.content; messages.push({ @@ -406,14 +405,12 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { }: ChatWrapperGenerateContextStateOptions): ChatWrapperGeneratedContextState & { transformedSystemMessagesToUserMessages: boolean } { - const { - contextText, stopGenerationTriggers, ignoreStartText, functionCall, transformedSystemMessagesToUserMessages - } = this._generateContextState({ + const {endJinjaMessagesWithUserMessage, ...contextState} = this._generateContextState({ chatHistory, availableFunctions, documentFunctionParams, endJinjaMessagesWithUserMessage: this._endJinjaMessagesWithUserMessage }); - return {contextText, stopGenerationTriggers, ignoreStartText, functionCall, transformedSystemMessagesToUserMessages}; + return contextState; } public override addAvailableFunctionsSystemMessageToHistory( @@ -683,13 +680,8 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { const lastJinjaItem = jinjaItems.at(-1); let eraseRenderedJinjaFromId: string | undefined; - if (endJinjaMessagesWithUserMessage && lastJinjaItem?.role === this.modelRoleName && - typeof lastJinjaItem.content === "string" && - lastJinjaItem.content.length > 0 && - ( - (lastJinjaItem as OpenAiChatAssistantMessage)["tool_calls"] == null || - (lastJinjaItem as OpenAiChatAssistantMessage)["tool_calls"]?.length === 0 - ) + if (endJinjaMessagesWithUserMessage && isPlainModelMessage(lastJinjaItem, this.modelRoleName) && + lastJinjaItem.content.length > 0 ) { eraseRenderedJinjaFromId = lastJinjaItem.content; jinjaItems.push({ @@ -698,24 +690,27 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { } as OpenAiChatMessage); } - const renderJinjaText = () => { + const renderJinja = (messages: OpenAiChatMessage[], options: Record) => ( + this._jinjaTemplate.render({ + ...( + this.additionalRenderParameters == null + ? {} + : structuredClone(this.additionalRenderParameters) + ), + messages, + ...removeUndefinedFields({tools}), + "bos_token": bosTokenId, + "eos_token": eosTokenId, + "eot_token": eotTokenId, + ...options + }) + ); + + // the render inputs (`jinjaItems`, `tools`, `eraseRenderedJinjaFromId`) don't change past this point + const renderedJinjaText = (() => { let res = tryMatrix({ options: [{}, {"add_generation_prompt": true}] - }, ({options}) => ( - this._jinjaTemplate.render({ - ...( - this.additionalRenderParameters == null - ? {} - : structuredClone(this.additionalRenderParameters) - ), - messages: jinjaItems, - ...removeUndefinedFields({tools}), - "bos_token": bosTokenId, - "eos_token": eosTokenId, - "eot_token": eotTokenId, - ...options - }) - )); + }, ({options}) => renderJinja(jinjaItems, options)); if (eraseRenderedJinjaFromId != null) { const eraseIndex = res.lastIndexOf(eraseRenderedJinjaFromId); @@ -724,6 +719,51 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { } return res; + })(); + + // When the last message is an empty model message (a generation prompt), + // some templates emit static text between the assistant header and the model's + // response (like a pre-filled `` tag) only inside their `add_generation_prompt` + // block. That text is not produced when rendering the empty model message as a + // regular message, so it would otherwise be dropped from the context. + // This extracts that static text so it can be appended to the context. + const extractGenerationPromptPrefill = (): string | undefined => { + // when the template never consults `add_generation_prompt`, it cannot emit text + // exclusive to the generation prompt, so the extra render can be skipped + if (!lastItemIsModelMessage || endJinjaMessagesWithUserMessage || + !this.template.includes("add_generation_prompt") + ) + return undefined; + + // the model message must be empty, since a non-empty one is a response prefix + // that the model continues from, not a generation prompt + const lastModelJinjaItem = jinjaItems.at(-1); + if (!isPlainModelMessage(lastModelJinjaItem, this.modelRoleName) || + idToContent.get(lastModelJinjaItem.content)?.toString() !== "" + ) + return undefined; + + let generationPromptRender: string; + try { + generationPromptRender = renderJinja(jinjaItems.slice(0, -1), {"add_generation_prompt": true}); + } catch (err) { + return undefined; + } + + const commonPrefixLength = findCommonStartLength(generationPromptRender, renderedJinjaText); + + // The two renders must diverge exactly where the empty model message content is, + // to ensure the extra text belongs to the model generation position + if (!renderedJinjaText.slice(commonPrefixLength).startsWith(lastModelJinjaItem.content)) + return undefined; + + const prefill = generationPromptRender.slice(commonPrefixLength); + + // The extracted text must be static template text and not contain any message content + if (prefill.length === 0 || includesText(prefill, [...idToContent.keys()], true)) + return undefined; + + return prefill; }; const validateThatAllMessageIdsAreUsed = (parts: ReturnType>) => { @@ -741,7 +781,7 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { }; const renderJinjaAndSplitIntoParts = () => { - const splitJinjaParts = splitText(renderJinjaText(), [...idToContent.keys()]); + const splitJinjaParts = splitText(renderedJinjaText, [...idToContent.keys()]); if (lastItemIsModelMessage) { let lastModelResponseIndex = -1; @@ -780,6 +820,29 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { const {splitJinjaParts, stopGenerationJinjaParts} = renderJinjaAndSplitIntoParts(); + // Preserve static text a template emits after the assistant header for a generation + // prompt (like a pre-filled `` tag) that would otherwise be dropped + let noPrefixTrigger: Extract< + NonNullable, {type: "segment"} + > | undefined; + const generationPromptPrefill = extractGenerationPromptPrefill(); + if (generationPromptPrefill != null) { + const segmentType = this._findSegmentTypeForGenerationPromptPrefill(generationPromptPrefill); + + if (segmentType != null) + // The pre-filled text opens a segment (like a thought segment), so inject it and + // open the segment during generation instead of adding it to the context as-is, + // to ensure the model's output is correctly attributed to that segment + noPrefixTrigger = { + type: "segment", + segmentType, + inject: LlamaText(new SpecialTokensText(generationPromptPrefill)) + }; + else + // Otherwise keep the pre-filled text as part of the context + splitJinjaParts.push(generationPromptPrefill); + } + const messageIdsLeftToProcess = new Set(messageIds); const contextText = LlamaText( splitJinjaParts.map((part) => { @@ -800,46 +863,87 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { if (messageIdsLeftToProcess.size !== 0) throw new Error("Some input messages are not present in the generated Jinja template output"); + const stopGenerationTriggers = [ + LlamaText(new SpecialToken("EOS")), + ...( + stopGenerationJinjaParts.length === 0 + ? [] + : [ + LlamaText( + stopGenerationJinjaParts.map((part) => { + if (typeof part === "string") + return new SpecialTokensText(part); + + const message = idToContent.get(part.separator); + + if (message == null) + throw new Error(`Message with id "${part.separator}" not found`); + + return message; + }) + ) + ] + ) + ]; + return { contextText, - ignoreStartText: !this.trimLeadingWhitespaceInResponses - ? [] - : [ - // ignore up to 4 leading spaces - ...Array(4).fill(0) - .map((_, index) => LlamaText(" ".repeat(index + 1))), - LlamaText("\t"), - LlamaText("\t\t"), - LlamaText("\t "), - LlamaText(" \t") - ], - stopGenerationTriggers: [ - LlamaText(new SpecialToken("EOS")), - ...( - stopGenerationJinjaParts.length === 0 - ? [] - : [ - LlamaText( - stopGenerationJinjaParts.map((part) => { - if (typeof part === "string") - return new SpecialTokensText(part); - - const message = idToContent.get(part.separator); - - if (message == null) - throw new Error(`Message with id "${part.separator}" not found`); - - return message; - }) - ) - ] - ) - ], + ...( + noPrefixTrigger != null + // `ignoreStartText` is unavailable alongside `noPrefixTrigger`, + // since the injected text is followed directly by the model's output + ? {noPrefixTrigger} + : { + ignoreStartText: !this.trimLeadingWhitespaceInResponses + ? [] + : [ + // ignore up to 4 leading spaces + ...Array(4).fill(0) + .map((_, index) => LlamaText(" ".repeat(index + 1))), + LlamaText("\t"), + LlamaText("\t\t"), + LlamaText("\t "), + LlamaText(" \t") + ] + } + ), + stopGenerationTriggers, transformedSystemMessagesToUserMessages, endJinjaMessagesWithUserMessage }; } + /** + * Find the segment type (like a thought segment) that a generation prompt prefill _opens_, + * so a pre-filled segment opening (like ``) can open the matching segment during generation. + * + * Only matches when the prefill leaves the segment open (it contains the prefix but not the suffix). + * A prefill that also closes the segment (like ``, used to suppress thoughts) is not + * matched, so it's kept as plain context text instead of leaving a segment open during generation. + * @internal + */ + private _findSegmentTypeForGenerationPromptPrefill(prefill: string): ChatModelSegmentType | undefined { + for (const segmentType of allSegmentTypes) { + const segmentDefinition = getChatWrapperSegmentDefinition(this.settings, segmentType); + if (segmentDefinition == null) + continue; + + const segmentPrefix = LlamaText(segmentDefinition.prefix).toString(); + if (segmentPrefix.length === 0 || !prefill.startsWith(segmentPrefix)) + continue; + + const segmentSuffix = segmentDefinition.suffix == null + ? "" + : LlamaText(segmentDefinition.suffix).toString(); + if (segmentSuffix.length > 0 && prefill.includes(segmentSuffix)) + continue; + + return segmentType; + } + + return undefined; + } + /** * Validate that this Jinja template can be rendered * @internal @@ -881,6 +985,21 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { } } +/** + * Whether the given Jinja message is a plain model message - a model message that carries only text content, + * without any tool calls attached to it. + */ +function isPlainModelMessage( + item: OpenAiChatMessage | undefined, modelRoleName: string +): item is OpenAiChatMessage & {content: string} { + return item != null && item.role === modelRoleName && + typeof item.content === "string" && + ( + (item as OpenAiChatAssistantMessage)["tool_calls"] == null || + (item as OpenAiChatAssistantMessage)["tool_calls"]?.length === 0 + ); +} + function resolveConvertUnsupportedSystemMessagesToUserMessagesOption( convertUnsupportedSystemMessagesToUserMessages?: JinjaTemplateChatWrapperOptions["convertUnsupportedSystemMessagesToUserMessages"] ): JinjaTemplateChatWrapperOptionsConvertMessageFormat | undefined { diff --git a/src/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.ts b/src/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.ts index 1d18e631..414f33d0 100644 --- a/src/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.ts +++ b/src/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.ts @@ -870,7 +870,7 @@ function extractWhitespacePrefixFromRevivedText(target: LlamaText) { }; } -function findCommonStartLength(text1: string, text2: string) { +export function findCommonStartLength(text1: string, text2: string) { let commonStartLength = 0; while (commonStartLength < text1.length && commonStartLength < text2.length) { if (text1[commonStartLength] !== text2[commonStartLength]) diff --git a/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts b/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts index 0658e6af..55d54b93 100644 --- a/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts +++ b/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts @@ -1534,6 +1534,91 @@ describe("JinjaTemplateChatWrapper", () => { }); }); + test("preserves generation prompt text emitted after the assistant header", () => { + const chatWrapper = new JinjaTemplateChatWrapper({ + template: "{%- for m in messages %}{{- '<|im_start|>' + m.role + '\\n' + m.content + '<|im_end|>\\n' }}{%- endfor %}" + + "{%- if add_generation_prompt %}{{- '<|im_start|>assistant\\n\\n' }}{%- endif %}" + }); + const {contextText, stopGenerationTriggers} = chatWrapper.generateContextState({ + chatHistory: [{ + type: "user", + text: "Hi" + }, { + type: "model", + response: [] + }] + }); + + // the `` tag pre-filled by the template's generation prompt must be preserved. + // This template has no ``, so no thought segment is configured and the pre-filled + // text is kept inline in the context. + expect(contextText.toString()).toBe("<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n\n"); + expect(contextText.toJSON()).toEqual([ + {type: "specialTokensText", value: "<|im_start|>user\n"}, + "Hi", + {type: "specialTokensText", value: "<|im_end|>\n<|im_start|>assistant\n\n"} + ]); + expect(stopGenerationTriggers.map((trigger) => trigger.toString())).toEqual(["EOS", "<|im_end|>\n"]); + }); + + // renders `` for the assistant, so a thought segment is detected, + // and pre-fills the given text in the generation prompt + const createThoughtTemplateChatWrapper = (generationPrompt: string) => new JinjaTemplateChatWrapper({ + template: "{%- for m in messages %}" + + "{%- if m.role == 'assistant' %}{{- '<|im_start|>assistant\\n' + m.content + '\\n<|im_end|>\\n' }}" + + "{%- else %}{{- '<|im_start|>' + m.role + '\\n' + m.content + '<|im_end|>\\n' }}{%- endif %}" + + "{%- endfor %}" + + "{%- if add_generation_prompt %}{{- '" + generationPrompt + "' }}{%- endif %}" + }); + + test("opens a segment for a pre-filled thought prefix in the generation prompt", () => { + const chatWrapper = createThoughtTemplateChatWrapper("<|im_start|>assistant\\n\\n"); + + expect(chatWrapper.settings.segments?.thought).toBeDefined(); + + const contextState = chatWrapper.generateContextState({ + chatHistory: [{ + type: "user", + text: "Hi" + }, { + type: "model", + response: [] + }] + }); + + // the pre-filled `` opens a thought segment rather than being kept inline, + // so the model's output is correctly attributed to the thought segment + expect(contextState.contextText.toString()).toBe("<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n"); + expect(contextState.noPrefixTrigger).toMatchObject({ + type: "segment", + segmentType: "thought" + }); + expect(contextState.noPrefixTrigger?.inject?.toString()).toBe("\n"); + }); + + test("does not open a segment for a pre-filled closed thought in the generation prompt", () => { + // pre-fills a complete, already-closed empty thought in the generation prompt + // (as used to suppress thoughts), which must not open a thought segment during generation + const chatWrapper = createThoughtTemplateChatWrapper("<|im_start|>assistant\\n\\n\\n\\n\\n"); + + expect(chatWrapper.settings.segments?.thought).toBeDefined(); + + const contextState = chatWrapper.generateContextState({ + chatHistory: [{ + type: "user", + text: "Hi" + }, { + type: "model", + response: [] + }] + }); + + // the closed thought is kept inline in the context and no segment is opened + expect(contextState.contextText.toString()) + .toBe("<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"); + expect(contextState.noPrefixTrigger).toBeUndefined(); + }); + test("Fails when messages are not present in the render output", () => { try { new JinjaTemplateChatWrapper({