From 561e880654aa4e1311b22577e9149962be0d8ba0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Fri, 10 May 2024 06:49:40 +0200 Subject: [PATCH] fix(security): Render all jinja templates in immutable sandbox (#1441) Chat templates are rendered with ImmutableSandboxedEnvironment in transformers so no need to do otherwise here. Co-authored-by: Andrei --- llama_cpp/llama_chat_format.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/llama_cpp/llama_chat_format.py b/llama_cpp/llama_chat_format.py index e0c59c9..d795b30 100644 --- a/llama_cpp/llama_chat_format.py +++ b/llama_cpp/llama_chat_format.py @@ -685,8 +685,7 @@ def hf_tokenizer_config_to_chat_formatter( assert isinstance(tokenizer_config["eos_token"], str) eos_token = tokenizer_config["eos_token"] - env = jinja2.Environment( - loader=jinja2.BaseLoader(), + env = ImmutableSandboxedEnvironment( trim_blocks=True, lstrip_blocks=True, ).from_string(chat_template) @@ -2603,7 +2602,10 @@ class Llava15ChatHandler: messages = [llama_types.ChatCompletionRequestSystemMessage(role="system", content=self.DEFAULT_SYSTEM_MESSAGE)] + messages image_urls = self.get_image_urls(messages) - template = jinja2.Template(self.CHAT_FORMAT) + template = ImmutableSandboxedEnvironment( + trim_blocks=True, + lstrip_blocks=True, + ).from_string(self.CHAT_FORMAT) text = template.render( messages=messages, add_generation_prompt=True, @@ -3242,8 +3244,7 @@ def chatml_function_calling( "{% endfor %}" "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}" ) - template_renderer = jinja2.Environment( - loader=jinja2.BaseLoader(), + template_renderer = ImmutableSandboxedEnvironment( autoescape=jinja2.select_autoescape(["html", "xml"]), undefined=jinja2.StrictUndefined, ).from_string(function_calling_template)