|
"""
|
||
|
|
llama_cpp/llama_jinja_format.py
|
||
|
|
"""
|
||
|
|
import dataclasses
|
||
|
|
from typing import Any, Callable, Dict, List, Optional, Protocol, Union
|
||
|
|
|
||
|
|
import jinja2
|
||
|
|
from jinja2 import Template
|
||
|
|
|
||
|
|
# NOTE: We sacrifice readability for usability.
|
||
|
|
# It will fail to work as expected if we attempt to format it in a readable way.
|
||
|
|
llama2_template = """{% for message in messages %}{% if message['role'] == 'user' %}[INST] {{ message['content'] }} [/INST]\n{% elif message['role'] == 'assistant' %}{{ message['content'] }}\n{% elif message['role'] == 'system' %}<<SYS>> {{ message['content'] }} <</SYS>>\n{% endif %}{% endfor %}"""
|
||
|
|
|
||
|
|
|
||
|
|
class MetaSingleton(type):
|
||
|
|
"""
|
||
|
|
Metaclass for implementing the Singleton pattern.
|
||
|
|
"""
|
||
|
|
|
||
|
|
_instances = {}
|
||
|
|
|
||
|
|
def __call__(cls, *args, **kwargs):
|
||
|
|
if cls not in cls._instances:
|
||
|
|
cls._instances[cls] = super(MetaSingleton, cls).__call__(*args, **kwargs)
|
||
|
|
return cls._instances[cls]
|
||
|
|
|
||
|
|
|
||
|
|
class Singleton(object, metaclass=MetaSingleton):
|
||
|
|
"""
|
||
|
|
Base class for implementing the Singleton pattern.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super(Singleton, self).__init__()
|
||
|
|
|
||
|
|
|
||
|
|
@dataclasses.dataclass
|
||
|
|
class ChatFormatterResponse:
|
||
|
|
prompt: str
|
||
|
|
stop: Optional[Union[str, List[str]]] = None
|
||
|
|
|
||
|
|
|
||
|
|
# Base Chat Formatter Protocol
|
||
|
|
class ChatFormatterInterface(Protocol):
|
||
|
|
def __init__(self, template: Optional[object] = None):
|
||
|
|
...
|
||
|
|
|
||
|
|
def __call__(
|
||
|
|
self,
|
||
|
|
messages: List[Dict[str, str]],
|
||
|
|
**kwargs,
|
||
|
|
) -> ChatFormatterResponse:
|
||
|
|
...
|
||
|
|
|
||
|
|
@property
|
||
|
|
def template(self) -> str:
|
||
|
|
...
|
||
|
|
|
||
|
|
|
||
|
|
class AutoChatFormatter(ChatFormatterInterface):
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
template: Optional[str] = None,
|
||
|
|
template_class: Optional[Template] = None,
|
||
|
|
):
|
||
|
|
if template is not None:
|
||
|
|
self._template = template
|
||
|
|
else:
|
||
|
|
self._template = llama2_template # default template
|
||
|
|
|
||
|
|
self._environment = jinja2.Environment(
|
||
|
|
loader=jinja2.BaseLoader(),
|
||
|
|
trim_blocks=True,
|
||
|
|
lstrip_blocks=True,
|
||
|
|
).from_string(
|
||
|
|
self._template,
|
||
|
|
template_class=template_class,
|
||
|
|
)
|
||
|
|
|
||
|
|
def __call__(
|
||
|
|
self,
|
||
|
|
messages: List[Dict[str, str]],
|
||
|
|
**kwargs: Any,
|
||
|
|
) -> ChatFormatterResponse:
|
||
|
|
formatted_sequence = self._environment.render(messages=messages, **kwargs)
|
||
|
|
return ChatFormatterResponse(prompt=formatted_sequence)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def template(self) -> str:
|
||
|
|
return self._template
|
||
|
|
|
||
|
|
|
||
|
|
class FormatterNotFoundException(Exception):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class ChatFormatterFactory(Singleton):
|
||
|
|
_chat_formatters: Dict[str, Callable[[], ChatFormatterInterface]] = {}
|
||
|
|
|
||
|
|
def register_formatter(
|
||
|
|
self,
|
||
|
|
name: str,
|
||
|
|
formatter_callable: Callable[[], ChatFormatterInterface],
|
||
|
|
overwrite=False,
|
||
|
|
):
|
||
|
|
if not overwrite and name in self._chat_formatters:
|
||
|
|
raise ValueError(
|
||
|
|
f"Formatter with name '{name}' is already registered. Use `overwrite=True` to overwrite it."
|
||
|
|
)
|
||
|
|
self._chat_formatters[name] = formatter_callable
|
||
|
|
|
||
|
|
def unregister_formatter(self, name: str):
|
||
|
|
if name in self._chat_formatters:
|
||
|
|
del self._chat_formatters[name]
|
||
|
|
else:
|
||
|
|
raise ValueError(f"No formatter registered under the name '{name}'.")
|
||
|
|
|
||
|
|
def get_formatter_by_name(self, name: str) -> ChatFormatterInterface:
|
||
|
|
try:
|
||
|
|
formatter_callable = self._chat_formatters[name]
|
||
|
|
return formatter_callable()
|
||
|
|
except KeyError:
|
||
|
|
raise FormatterNotFoundException(
|
||
|
|
f"Invalid chat format: {name} (valid formats: {list(self._chat_formatters.keys())})"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
# Define a chat format class
|
||
|
|
class Llama2Formatter(AutoChatFormatter):
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__(llama2_template)
|
||
|
|
|
||
|
|
|
||
|
|
# With the Singleton pattern applied, regardless of where or how many times
|
||
|
|
# ChatFormatterFactory() is called, it will always return the same instance
|
||
|
|
# of the factory, ensuring that the factory's state is consistent throughout
|
||
|
|
# the application.
|
||
|
|
ChatFormatterFactory().register_formatter("llama-2", Llama2Formatter)
|