Add Pygmalion chat format (#986)
This commit is contained in:
parent
6bbeea07ae
commit
b938cccf05
1 changed files with 17 additions and 0 deletions
|
@ -637,6 +637,23 @@ def format_zephyr(
|
||||||
_prompt = _format_chatml(system_message, _messages, _sep)
|
_prompt = _format_chatml(system_message, _messages, _sep)
|
||||||
return ChatFormatterResponse(prompt=_prompt, stop=_sep)
|
return ChatFormatterResponse(prompt=_prompt, stop=_sep)
|
||||||
|
|
||||||
|
|
||||||
|
@register_chat_format("pygmalion")
|
||||||
|
def format_pygmalion(
|
||||||
|
messages: List[llama_types.ChatCompletionRequestMessage],
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> ChatFormatterResponse:
|
||||||
|
system_template = """<|system|>{system_message}"""
|
||||||
|
system_message = _get_system_message(messages)
|
||||||
|
system_message = system_template.format(system_message=system_message)
|
||||||
|
_roles = dict(user="<|user|>", assistant="<|model|>")
|
||||||
|
_sep = "\n"
|
||||||
|
_messages = _map_roles(messages, _roles)
|
||||||
|
_messages.append((_roles["assistant"], None))
|
||||||
|
_prompt = _format_chatml(system_message, _messages, _sep)
|
||||||
|
return ChatFormatterResponse(prompt=_prompt, stop=_sep)
|
||||||
|
|
||||||
|
|
||||||
@register_chat_format("chatml")
|
@register_chat_format("chatml")
|
||||||
def format_chatml(
|
def format_chatml(
|
||||||
messages: List[llama_types.ChatCompletionRequestMessage],
|
messages: List[llama_types.ChatCompletionRequestMessage],
|
||||||
|
|
Loading…
Reference in a new issue