Skip to content

Create tags attribute to fix MarkupError in rich CLI #2553

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merrymercy merged 1 commit into from
Oct 13, 2023
Merged

Create tags attribute to fix MarkupError in rich CLI #2553

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 10 additions & 6 deletions fastchat/conversation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

import dataclasses
from enum import auto, IntEnum
from typing import List, Any, Dict, Union
from typing import List, Any, Dict, Union, Tuple


class SeparatorStyle(IntEnum):
Expand Down Expand Up @@ -41,7 +41,7 @@ class Conversation:
# The system message
system_message: str = ""
# The names of two roles
roles: List[str] = (("USER", "ASSISTANT"),)
roles: Tuple[str] = ("USER", "ASSISTANT")
# All messages. Each item is (role, message).
messages: List[List[str]] = ()
# The number of few shot examples
Expand All @@ -54,6 +54,8 @@ class Conversation:
stop_str: Union[str, List[str]] = None
# Stops generation if meeting any token in this list
stop_token_ids: List[int] = None
# Tags to be used in the template
tags: Tuple[str] = None

def get_prompt(self) -> str:
"""Get the prompt for generation."""
Expand Down Expand Up @@ -128,13 +130,14 @@ def get_prompt(self) -> str:
else:
ret = "[INST] "
for i, (role, message) in enumerate(self.messages):
tag = self.tags[i % 2]
if message:
if i == 0:
ret += message + " "
else:
ret += role + " " + message + seps[i % 2]
ret += tag + " " + message + seps[i % 2]
else:
ret += role
ret += tag
return ret
elif self.sep_style == SeparatorStyle.CHATGLM:
# source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
Expand Down Expand Up @@ -266,6 +269,7 @@ def copy(self):
sep2=self.sep2,
stop_str=self.stop_str,
stop_token_ids=self.stop_token_ids,
tags=self.tags,
)

def dict(self):
Expand Down Expand Up @@ -846,7 +850,7 @@ def get_conv_template(name: str) -> Conversation:
Conversation(
name="mistral",
system_template="",
roles=("[INST]", "[/INST]"),
tags=("[INST]", "[/INST]"),
sep_style=SeparatorStyle.LLAMA2,
sep=" ",
sep2="</s>",
Expand All @@ -860,7 +864,7 @@ def get_conv_template(name: str) -> Conversation:
Conversation(
name="llama-2",
system_template="[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n",
roles=("[INST]", "[/INST]"),
tags=("[INST]", "[/INST]"),
sep_style=SeparatorStyle.LLAMA2,
sep=" ",
sep2=" </s><s>",
Expand Down