diff --git a/pilot/conversation.py b/pilot/conversation.py new file mode 100644 index 000000000..4db5d9548 --- /dev/null +++ b/pilot/conversation.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python3 +# -*- coding:utf-8 -*- + +import dataclasses +from enum import auto, Enum +from typing import List, Tuple, Any + + +class SeparatorStyle(Enum): + + SINGLE = auto() + TWO = auto() + +@dataclasses.dataclass +class Conversation: + """This class keeps all conversation history. """ + + system: str + roles: List[str] + messages: List[List[str]] + offset: int + sep_style: SeparatorStyle = SeparatorStyle.SINGLE + sep: str = "###" + sep2: str = None + + # Used for gradio server + skip_next: bool = False + conv_id: Any = None + + def get_prompt(self): + if self.sep_style == SeparatorStyle.SINGLE: + ret = self.system + for role, message in self.messages: + if message: + ret += self.sep + " " + role + ": " + message + else: + ret += self.sep + " " + role + ":" + return ret + + elif self.sep_style == SeparatorStyle.TWO: + seps = [self.sep, self.sep2] + ret = self.system + seps[0] + 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, + 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, + "roles": self.roles, + "messages": self.messages, + "offset": self.offset, + "sep": self.sep, + "sep2": self.sep2, + "conv_id": self.conv_id + } + + +conv_one_shot = Conversation( + system="A chat between a curious human and an artificial intelligence assistant, who very familiar with database related knowledge. " + "The assistant gives helpful, detailed, professional and polite answers to the human's questions. ", + roles=("Human", "Assistant"), + messages=( + ( + "Human", + "What are the key differences between mysql and postgres?", + ), + ( + "Assistant", + "MySQL and PostgreSQL are both popular open-source relational database management systems (RDBMS) " + "that have many similarities but also some differences. Here are some key differences: \n" + "1. Data Types: PostgreSQL has a more extensive set of data types, " + "including support for array, hstore, JSON, and XML, whereas MySQL has a more limited set.\n" + "2. ACID compliance: Both MySQL and PostgreSQL support ACID compliance (Atomicity, Consistency, Isolation, Durability), " + "but PostgreSQL is generally considered to be more strict in enforcing it.\n" + "3. Replication: MySQL has a built-in replication feature, which allows you to replicate data across multiple servers," + "whereas PostgreSQL has a similar feature, but it is not as mature as MySQL's.\n" + "4. Performance: MySQL is generally considered to be faster and more efficient in handling large datasets, " + "whereas PostgreSQL is known for its robustness and reliability.\n" + "5. Licensing: MySQL is licensed under the GPL (General Public License), which means that it is free and open-source software, " + "whereas PostgreSQL is licensed under the PostgreSQL License, which is also free and open-source but with different terms.\n" + + "Ultimately, the choice between MySQL and PostgreSQL depends on the specific needs and requirements of your application. " + "Both are excellent database management systems, and choosing the right one " + "for your project requires careful consideration of your application's requirements, performance needs, and scalability." + ), + ), + offset=2, + sep_style=SeparatorStyle.SINGLE, + sep="###" +) + +conv_vicuna_v1 = Conversation( + system = "A chat between a curious user and an artificial intelligence assistant, who very familiar with database related knowledge. " + "The assistant gives helpful, detailed, professional and polite answers to the user's questions. ", + roles=("USER", "ASSISTANT"), + messages=(), + offset=0, + sep_style=SeparatorStyle.TWO, + sep=" ", + sep2="", +) + +conv_template = { + "conv_one_shot": conv_one_shot, + "vicuna_v1": conv_vicuna_v1 +} + + +def get_default_conv_template(model_name: str = "vicuna-13b"): + model_name = model_name.lower() + if "vicuna" in model_name: + return conv_vicuna_v1 + return conv_one_shot + + +def compute_skip_echo_len(prompt): + skip_echo_len = len(prompt) + 1 - prompt.count("") * 3 + return skip_echo_len \ No newline at end of file