Add support for Jinja for chat model templates (#1739)

* Add support for Jinja for chat model templates

* feat(docs): improve docs regarding templates

---------

Co-authored-by: Nathan Sarrazin <[email protected]>

PROMPTS.md

@@ -1,5 +1,8 @@
 # Prompt templates
 
+> [!WARNING]
+> We now recommend using the `tokenizer` field to get the chat template directly from the hub. Just set it to your model id on the hub to automatically get the template.
+
 These are the templates used to format the conversation history for different models used in HuggingChat. Set them in your `.env.local` [like so](https://github.com/huggingface/chat-ui#chatprompttemplate).
 
 ## Llama 2

README.md

@@ -314,6 +314,9 @@ The following is the default `chatPromptTemplate`, although newlines and indenti
 {{assistantMessageToken}}
 ```
 
+> [!INFO]
+> We also support Jinja2 templates for the `chatPromptTemplate` in addition to Handlebars templates. On startup we first try to compile with Jinja and if that fails we fall back to interpreting `chatPromptTemplate` as handlebars.
+
 #### Multi modal model
 
 We currently support [IDEFICS](https://huggingface.co/blog/idefics) (hosted on TGI), OpenAI and Claude 3 as multimodal models. You can enable it by setting `multimodal: true` in your `MODELS` configuration. For IDEFICS, you must have a [PRO HF Api token](https://huggingface.co/settings/tokens). For OpenAI, see the [OpenAI section](#openai-api-compatible-models). For Anthropic, see the [Anthropic section](#anthropic).

src/lib/utils/template.spec.ts

@@ -0,0 +1,59 @@
+import { describe, test, expect } from "vitest";
+import { compileTemplate } from "./template";
+
+// Test data for simple templates
+const modelData = {
+	preprompt: "Hello",
+};
+
+const simpleTemplate = "Test: {{preprompt}} and {{foo}}";
+
+// Additional realistic test data for Llama 70B templates
+const messages = [
+	{ from: "user", content: "Hello there" },
+	{ from: "assistant", content: "Hi, how can I help?" },
+];
+
+// Handlebars Llama 70B Template
+const llama70bTemplateHB = `<s>{{#if preprompt}}Source: system\n\n{{preprompt}}<step>{{/if}}{{#each messages}}{{#ifUser}}Source: user\n\n{{content}}<step>{{/ifUser}}{{#ifAssistant}}Source: assistant\n\n{{content}}<step>{{/ifAssistant}}{{/each}}Source: assistant\nDestination: user\n\n`;
+
+// Expected output for Handlebars Llama 70B Template
+const expectedHB =
+	"<s>Source: system\n\nSystem Message<step>Source: user\n\nHello there<step>Source: assistant\n\nHi, how can I help?<step>Source: assistant\nDestination: user\n\n";
+
+// Jinja Llama 70B Template
+const llama70bTemplateJinja = `<s>{% if preprompt %}Source: system\n\n{{ preprompt }}<step>{% endif %}{% for message in messages %}{% if message.from == 'user' %}Source: user\n\n{{ message.content }}<step>{% elif message.from == 'assistant' %}Source: assistant\n\n{{ message.content }}<step>{% endif %}{% endfor %}Source: assistant\nDestination: user\n\n`;
+
+// Expected output for Jinja Llama 70B Template
+const expectedJinja =
+	"<s>Source: system\n\nSystem Message<step>Source: user\n\nHello there<step>Source: assistant\n\nHi, how can I help?<step>Source: assistant\nDestination: user\n\n";
+
+describe("Template Engine Rendering", () => {
+	test("should render using Handlebars fallback when no templateEngine is specified", () => {
+		const render = compileTemplate(simpleTemplate, modelData);
+		const result = render({ foo: "World" });
+		expect(result).toBe("Test: Hello and World");
+	});
+
+	test('should render using Jinja when templateEngine is set to "jinja"', () => {
+		const render = compileTemplate(simpleTemplate, modelData);
+		const result = render({ foo: "World" });
+		expect(result).toBe("Test: Hello and World");
+	});
+
+	// Realistic Llama 70B template tests
+	test("should render realistic Llama 70B template using Handlebars", () => {
+		const render = compileTemplate(llama70bTemplateHB, { preprompt: "System Message" });
+		const result = render({ messages });
+		expect(result).toBe(expectedHB);
+	});
+
+	test("should render realistic Llama 70B template using Jinja", () => {
+		const render = compileTemplate(llama70bTemplateJinja, {
+			preprompt: "System Message",
+		});
+		const result = render({ messages });
+		// Trim both outputs to account for whitespace differences in Jinja engine
+		expect(result.trim()).toBe(expectedJinja.trim());
+	});
+});

src/lib/utils/template.ts

@@ -1,6 +1,8 @@
 import type { Message } from "$lib/types/Message";
 import Handlebars from "handlebars";
+import { Template } from "@huggingface/jinja";
 
+// Register Handlebars helpers
 Handlebars.registerHelper("ifUser", function (this: Pick<Message, "from" | "content">, options) {
 	if (this.from == "user") return options.fn(this);
 });
@@ -12,8 +14,21 @@ Handlebars.registerHelper(
 	}
 );
 
-export function compileTemplate<T>(input: string, model: { preprompt: string }) {
-	const template = Handlebars.compile<T>(input, {
+// Updated compileTemplate to try Jinja and fallback to Handlebars if Jinja fails
+export function compileTemplate<T>(
+	input: string,
+	model: { preprompt: string; templateEngine?: string }
+) {
+	let jinjaTemplate: Template | undefined;
+	try {
+		// Try to compile with Jinja
+		jinjaTemplate = new Template(input);
+	} catch (e) {
+		// Could not compile with Jinja
+		jinjaTemplate = undefined;
+	}
+
+	const hbTemplate = Handlebars.compile<T>(input, {
 		knownHelpers: { ifUser: true, ifAssistant: true },
 		knownHelpersOnly: true,
 		noEscape: true,
@@ -21,7 +36,15 @@ export function compileTemplate<T>(input: string, model: { preprompt: string })
 		preventIndent: true,
 	});
 
-	return function render(inputs: T, options?: RuntimeOptions) {
-		return template({ ...model, ...inputs }, options);
+	return function render(inputs: T) {
+		if (jinjaTemplate) {
+			try {
+				return jinjaTemplate.render({ ...model, ...inputs });
+			} catch (e) {
+				// Fallback to Handlebars if Jinja rendering fails
+				return hbTemplate({ ...model, ...inputs });
+			}
+		}
+		return hbTemplate({ ...model, ...inputs });
 	};
 }