Add Pygmalion chat format (#986)

This commit is contained in:
chiensen 2023-12-12 09:44:04 +08:00 committed by GitHub
parent 6bbeea07ae
commit b938cccf05
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -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],