mirror of
https://github.com/theroyallab/tabbyAPI.git
synced 2026-03-14 15:57:27 +00:00
de9a19b5d3
Append generation prompts if given the flag on an OAI chat completion request. This appends the "assistant" message to the instruct prompt. Defaults to true since this is intended behavior. Signed-off-by: kingbri <bdashore3@proton.me>
42 lines
1.6 KiB
Python
42 lines
1.6 KiB
Python
import pathlib
|
|
from functools import lru_cache
|
|
from importlib.metadata import version as package_version
|
|
from jinja2.sandbox import ImmutableSandboxedEnvironment
|
|
from packaging import version
|
|
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, add_generation_prompt: bool):
|
|
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,
|
|
add_generation_prompt = add_generation_prompt
|
|
)
|
|
|
|
# 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
|
|
|
|
def get_template_from_file(prompt_template_name: str):
|
|
with open(pathlib.Path(f"templates/{prompt_template_name}.jinja"), "r") as raw_template:
|
|
return PromptTemplate(
|
|
name = prompt_template_name,
|
|
template = raw_template.read()
|
|
)
|