156 lines
4.6 KiB
Python
156 lines
4.6 KiB
Python
import dataclasses
|
|
from enum import auto, Enum
|
|
from typing import List, Tuple, Any
|
|
|
|
|
|
class SeparatorStyle(Enum):
|
|
"""Different separator style."""
|
|
SINGLE = auto()
|
|
TWO = auto()
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Conversation:
|
|
"""A class that keeps all conversation history."""
|
|
system: str
|
|
instruction: str
|
|
roles: List[str]
|
|
messages: List[List[str]]
|
|
offset: int
|
|
sep_style: SeparatorStyle = SeparatorStyle.SINGLE
|
|
sep: str = "###"
|
|
sep2: str = None
|
|
|
|
skip_next: bool = False
|
|
conv_id: Any = None
|
|
|
|
def get_prompt(self):
|
|
if self.sep_style == SeparatorStyle.SINGLE:
|
|
ret = self.system + self.sep
|
|
if self.instruction is not None and len(self.instruction) > 0:
|
|
ret += self.roles[2] + ": " + self.instruction + self.sep
|
|
for role, message in self.messages:
|
|
if message:
|
|
ret += role + ": " + message + self.sep
|
|
else:
|
|
ret += role + ":"
|
|
return ret
|
|
elif self.sep_style == SeparatorStyle.TWO:
|
|
seps = [self.sep, self.sep2]
|
|
ret = self.system + seps[0]
|
|
if self.instruction is not None and len(self.instruction) > 0:
|
|
ret += self.roles[2] + ": " + self.instruction + self.sep
|
|
for i, (role, message) in enumerate(self.messages):
|
|
if message:
|
|
ret += role + ": " + message + seps[i % 2]
|
|
else:
|
|
ret += role + ":"
|
|
return ret
|
|
else:
|
|
raise ValueError(f"Invalid style: {self.sep_style}")
|
|
|
|
def append_message(self, role, message):
|
|
self.messages.append([role, message])
|
|
|
|
def to_gradio_chatbot(self):
|
|
ret = []
|
|
for i, (role, msg) in enumerate(self.messages[self.offset:]):
|
|
if i % 2 == 0:
|
|
ret.append([msg, None])
|
|
else:
|
|
ret[-1][-1] = msg
|
|
return ret
|
|
|
|
def copy(self):
|
|
return Conversation(
|
|
system=self.system,
|
|
instruction=self.instruction,
|
|
roles=self.roles,
|
|
messages=[[x, y] for x, y in self.messages],
|
|
offset=self.offset,
|
|
sep_style=self.sep_style,
|
|
sep=self.sep,
|
|
sep2=self.sep2,
|
|
conv_id=self.conv_id)
|
|
|
|
def dict(self):
|
|
return {
|
|
"system": self.system,
|
|
"instruction": self.instruction,
|
|
"roles": self.roles,
|
|
"messages": self.messages,
|
|
"offset": self.offset,
|
|
"sep": self.sep,
|
|
"sep2": self.sep2,
|
|
"conv_id": self.conv_id,
|
|
}
|
|
|
|
|
|
conv_v1 = Conversation(
|
|
system="A chat between a curious human and an artificial intelligence assistant. "
|
|
"The assistant gives helpful, detailed, and polite answers to the human's questions.",
|
|
instruction="",
|
|
roles=("Human", "Assistant", "System"),
|
|
messages=(),
|
|
offset=0,
|
|
sep_style=SeparatorStyle.SINGLE,
|
|
sep="###",
|
|
)
|
|
|
|
conv_v1_2 = Conversation(
|
|
system="A chat between a curious human and an artificial intelligence assistant. "
|
|
"The assistant gives helpful, detailed, and polite answers to the human's questions.",
|
|
instruction="",
|
|
roles=("Human", "Assistant", "System"),
|
|
messages=(),
|
|
offset=0,
|
|
sep_style=SeparatorStyle.SINGLE,
|
|
sep="###",
|
|
)
|
|
|
|
conv_bair_v1 = Conversation(
|
|
system="BEGINNING OF CONVERSATION:",
|
|
instruction="",
|
|
roles=("USER", "GPT", "System"),
|
|
messages=(),
|
|
offset=0,
|
|
sep_style=SeparatorStyle.TWO,
|
|
sep=" ",
|
|
sep2="</s>",
|
|
)
|
|
|
|
|
|
default_conversation = conv_v1_2
|
|
conv_templates = {
|
|
"v1": conv_v1_2,
|
|
"bair_v1": conv_bair_v1,
|
|
}
|
|
|
|
def covert_prompt_to_input_ids_with_history(text, history, tokenizer, max_token):
|
|
conv = default_conversation.copy()
|
|
|
|
conv.append_message(conv.roles[1], None)
|
|
conv.append_message(conv.roles[0], text)
|
|
|
|
example = tokenizer.encode_plus(f"{conv.get_prompt()}", None, max_length=None)['input_ids']
|
|
|
|
while(len(history) > 0 and (len(example) < max_token)):
|
|
tmp = history.pop()
|
|
if tmp[0] == 'ASSISTANT':
|
|
conv.append_message(conv.roles[1], tmp[1])
|
|
else:
|
|
conv.append_message(conv.roles[0], tmp[1])
|
|
example = tokenizer.encode_plus(f"{conv.get_prompt()}", None, max_length=None)['input_ids']
|
|
|
|
if len(example) >= max_token:
|
|
conv.messages.pop()
|
|
conv.messages = conv.messages[::-1]
|
|
print('model in:', conv.get_prompt())
|
|
example = tokenizer.encode_plus(f"{conv.get_prompt()}", None, max_length=None)['input_ids']
|
|
example = example[1:-1]
|
|
|
|
return example
|
|
|
|
if __name__ == "__main__":
|
|
print(default_conversation.get_prompt())
|