Skip to content

Commit

Permalink
Add baichuan-2 chat format (abetlen#936)
Browse files Browse the repository at this point in the history
Signed-off-by: caiyesd <[email protected]>
  • Loading branch information
caiyesd authored Nov 22, 2023
1 parent 9515467 commit b8f29f4
Showing 1 changed file with 15 additions and 0 deletions.
15 changes: 15 additions & 0 deletions llama_cpp/llama_chat_format.py
Original file line number Diff line number Diff line change
Expand Up @@ -456,6 +456,21 @@ def format_oasst_llama(
return ChatFormatterResponse(prompt=_prompt)


@register_chat_format("baichuan-2")
def format_baichuan2(
messages: List[llama_types.ChatCompletionRequestMessage],
**kwargs: Any,
) -> ChatFormatterResponse:
_system_template = "{system_message}"
_roles = dict(user="<reserved_106>", assistant="<reserved_107>")
_sep = ""
system_message = _get_system_message(messages)
system_message = _system_template.format(system_message=system_message)
_messages = _map_roles(messages, _roles)
_messages.append((_roles["assistant"], None))
_prompt = _format_no_colon_single(system_message, _messages, _sep)
return ChatFormatterResponse(prompt=_prompt)

@register_chat_format("openbuddy")
def format_openbuddy(
messages: List[llama_types.ChatCompletionRequestMessage],
Expand Down

0 comments on commit b8f29f4

Please sign in to comment.