mirror of
https://github.com/theroyallab/tabbyAPI.git
synced 2026-03-15 00:07:28 +00:00
f631dd6ff7
Jinja2 is a lightweight template parser that's used in Transformers for parsing chat completions. It's much more efficient than Fastchat and can be imported as part of requirements. Also allows for unblocking Pydantic's version. Users now have to provide their own template if needed. A separate repo may be usable for common prompt template storage. Signed-off-by: kingbri <bdashore3@proton.me>
31 lines
1.2 KiB
Python
31 lines
1.2 KiB
Python
from functools import lru_cache
|
|
from importlib.metadata import version as package_version
|
|
from packaging import version
|
|
from jinja2.sandbox import ImmutableSandboxedEnvironment
|
|
from pydantic import BaseModel
|
|
|
|
# Small replication of AutoTokenizer's chat template system for efficiency
|
|
|
|
class PromptTemplate(BaseModel):
|
|
name: str
|
|
template: str
|
|
|
|
def get_prompt_from_template(messages, prompt_template: PromptTemplate):
|
|
if version.parse(package_version("jinja2")) < version.parse("3.0.0"):
|
|
raise ImportError(
|
|
"Parsing these chat completion messages requires fastchat 0.2.23 or greater. "
|
|
f"Current version: {version('jinja2')}\n"
|
|
"Please upgrade fastchat by running the following command: "
|
|
"pip install -U fschat[model_worker]"
|
|
)
|
|
|
|
compiled_template = _compile_template(prompt_template.template)
|
|
return compiled_template.render(messages = messages)
|
|
|
|
# Inspired from https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L1761
|
|
@lru_cache
|
|
def _compile_template(template: str):
|
|
jinja_env = ImmutableSandboxedEnvironment(trim_blocks = True, lstrip_blocks = True)
|
|
jinja_template = jinja_env.from_string(template)
|
|
return jinja_template
|