Skip to content

Commit

Permalink
fix: use both eos and bos tokens as stop sequences for hf-tokenizer-c…
Browse files Browse the repository at this point in the history
…onfig chat format.
  • Loading branch information
abetlen committed Jan 22, 2024
1 parent 2ce0b8a commit 5b982d0
Showing 1 changed file with 12 additions and 8 deletions.
20 changes: 12 additions & 8 deletions llama_cpp/llama_chat_format.py
Original file line number Diff line number Diff line change
Expand Up @@ -379,7 +379,8 @@ def hf_autotokenizer_to_chat_completion_handler(


def hf_tokenizer_config_to_chat_formatter(
tokenizer_config: Dict[str, Any]
tokenizer_config: Dict[str, Any],
add_generation_prompt: bool = True,
) -> ChatFormatter:
assert isinstance(tokenizer_config, dict)

Expand All @@ -401,31 +402,34 @@ def hf_tokenizer_config_to_chat_formatter(
lstrip_blocks=True,
).from_string(chat_template)

def format_autotokenizer(
def format_tokenizer_config(
messages: List[llama_types.ChatCompletionRequestMessage],
**kwargs: Any,
) -> ChatFormatterResponse:
# TODO: veryify this is correct
# Add a blank assistant message to the end of the messages to prompt the model to generate a response
prompt = env.render(
messages=[
if add_generation_prompt:
messages = [
*messages,
llama_types.ChatCompletionRequestAssistantMessage(
role="assistant", content=""
),
],
]
prompt = env.render(
messages=messages,
bos_token=bos_token,
eos_token=eos_token,
)
return ChatFormatterResponse(prompt=prompt, stop=eos_token)
return ChatFormatterResponse(prompt=prompt, stop=[eos_token, bos_token])

return format_autotokenizer
return format_tokenizer_config


def hf_tokenizer_config_to_chat_completion_handler(
tokenizer_config: Dict[str, Any],
add_generation_prompt: bool = True,
) -> LlamaChatCompletionHandler:
chat_formatter = hf_tokenizer_config_to_chat_formatter(tokenizer_config)
chat_formatter = hf_tokenizer_config_to_chat_formatter(tokenizer_config, add_generation_prompt=add_generation_prompt)
return chat_formatter_to_chat_completion_handler(chat_formatter)


Expand Down

0 comments on commit 5b982d0

Please sign in to comment.