baalajimaestro/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

fix(security): Render all jinja templates in immutable sandbox (#1441)

Browse source 
Chat templates are rendered with ImmutableSandboxedEnvironment in transformers so no need to do otherwise here.

Co-authored-by: Andrei <abetlen@gmail.com>
This commit is contained in:
Sigbjørn Skjæret 2024-05-10 06:49:40 +02:00 committed by GitHub
parent b454f40a9a
commit 561e880654
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 6 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
llama_cpp/llama_chat_format.py
View file

@ -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)