Singularity666/editx
1
1import dataclasses2from enum import auto, Enum3from typing import List, Tuple4 5 6class SeparatorStyle(Enum):7 """Different separator style."""8 SINGLE = auto()9 TWO = auto()10 MPT = auto()11 12 13@dataclasses.dataclass14class Conversation:15 """A class that keeps all conversation history."""16 system: str17 roles: List[str]18 messages: List[List[str]]19 offset: int20 sep_style: SeparatorStyle = SeparatorStyle.SINGLE21 sep: str = "###"22 sep2: str = None23 version: str = "Unknown"24 25 skip_next: bool = False26 27 def get_prompt(self):28 if self.sep_style == SeparatorStyle.SINGLE:29 ret = self.system + self.sep30 for role, message in self.messages:31 if message:32 if type(message) is tuple:33 message, _, _ = message34 ret += role + ": " + message + self.sep35 else:36 ret += role + ":"37 return ret38 elif self.sep_style == SeparatorStyle.TWO:39 seps = [self.sep, self.sep2]40 ret = self.system + seps[0]41 for i, (role, message) in enumerate(self.messages):42 if message:43 if type(message) is tuple:44 message, _, _ = message45 ret += role + ": " + message + seps[i % 2]46 else:47 ret += role + ":"48 return ret49 if self.sep_style == SeparatorStyle.MPT:50 ret = self.system + self.sep51 for role, message in self.messages:52 if message:53 if type(message) is tuple:54 message, _, _ = message55 ret += role + message + self.sep56 else:57 ret += role58 return ret59 else:60 raise ValueError(f"Invalid style: {self.sep_style}")61 62 def append_message(self, role, message):63 self.messages.append([role, message])64 65 def get_images(self, return_pil=False):66 images = []67 for i, (role, msg) in enumerate(self.messages[self.offset:]):68 if i % 2 == 0:69 if type(msg) is tuple:70 import base6471 from io import BytesIO72 from PIL import Image73 msg, image, image_process_mode = msg74 if image_process_mode == "Pad":75 def expand2square(pil_img, background_color=(122, 116, 104)):76 width, height = pil_img.size77 if width == height:78 return pil_img79 elif width > height:80 result = Image.new(pil_img.mode, (width, width), background_color)81 result.paste(pil_img, (0, (width - height) // 2))82 return result83 else:84 result = Image.new(pil_img.mode, (height, height), background_color)85 result.paste(pil_img, ((height - width) // 2, 0))86 return result87 image = expand2square(image)88 elif image_process_mode == "Crop":89 pass90 elif image_process_mode == "Resize":91 image = image.resize((224, 224))92 else:93 raise ValueError(f"Invalid image_process_mode: {image_process_mode}")94 max_hw, min_hw = max(image.size), min(image.size)95 aspect_ratio = max_hw / min_hw96 max_len, min_len = 800, 40097 shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))98 longest_edge = int(shortest_edge * aspect_ratio)99 W, H = image.size100 if H > W:101 H, W = longest_edge, shortest_edge102 else:103 H, W = shortest_edge, longest_edge104 image = image.resize((W, H))105 if return_pil:106 images.append(image)107 else:108 buffered = BytesIO()109 image.save(buffered, format="JPEG")110 img_b64_str = base64.b64encode(buffered.getvalue()).decode()111 images.append(img_b64_str)112 return images113 114 def to_gradio_chatbot(self):115 ret = []116 for i, (role, msg) in enumerate(self.messages[self.offset:]):117 if i % 2 == 0:118 if type(msg) is tuple:119 import base64120 from io import BytesIO121 msg, image, image_process_mode = msg122 max_hw, min_hw = max(image.size), min(image.size)123 aspect_ratio = max_hw / min_hw124 max_len, min_len = 800, 400125 shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))126 longest_edge = int(shortest_edge * aspect_ratio)127 W, H = image.size128 if H > W:129 H, W = longest_edge, shortest_edge130 else:131 H, W = shortest_edge, longest_edge132 image = image.resize((W, H))133 # image = image.resize((224, 224))134 buffered = BytesIO()135 image.save(buffered, format="JPEG")136 img_b64_str = base64.b64encode(buffered.getvalue()).decode()137 img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'138 msg = msg.replace('<image>', img_str)139 ret.append([msg, None])140 else:141 ret[-1][-1] = msg142 return ret143 144 def copy(self):145 return Conversation(146 system=self.system,147 roles=self.roles,148 messages=[[x, y] for x, y in self.messages],149 offset=self.offset,150 sep_style=self.sep_style,151 sep=self.sep,152 sep2=self.sep2)153 154 def dict(self):155 if len(self.get_images()) > 0:156 return {157 "system": self.system,158 "roles": self.roles,159 "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],160 "offset": self.offset,161 "sep": self.sep,162 "sep2": self.sep2,163 }164 return {165 "system": self.system,166 "roles": self.roles,167 "messages": self.messages,168 "offset": self.offset,169 "sep": self.sep,170 "sep2": self.sep2,171 }172 173 174conv_v1 = Conversation(175 system="A chat between a curious human and an artificial intelligence assistant. "176 "The assistant gives helpful, detailed, and polite answers to the human's questions.",177 roles=("Human", "Assistant"),178 messages=(179 ("Human", "Give three tips for staying healthy."),180 ("Assistant",181 "Sure, here are three tips for staying healthy:\n"182 "1. Exercise regularly: Regular physical activity can help improve your overall health and wellbeing. "183 "It can also help reduce your risk of chronic conditions such as obesity, diabetes, heart disease, "184 "and certain cancers. Aim for at least 150 minutes of moderate-intensity aerobic exercise or "185 "75 minutes of vigorous-intensity aerobic exercise per week, along with muscle-strengthening "186 "activities at least two days per week.\n"187 "2. Eat a balanced diet: Eating a balanced diet that is rich in fruits, "188 "vegetables, whole grains, lean proteins, and healthy fats can help support "189 "your overall health. Try to limit your intake of processed and high-sugar foods, "190 "and aim to drink plenty of water throughout the day.\n"191 "3. Get enough sleep: Getting enough quality sleep is essential for your physical "192 "and mental health. Adults should aim for seven to nine hours of sleep per night. "193 "Establish a regular sleep schedule and try to create a relaxing bedtime routine to "194 "help improve the quality of your sleep.")195 ),196 offset=2,197 sep_style=SeparatorStyle.SINGLE,198 sep="###",199)200 201conv_v1_2 = Conversation(202 system="A chat between a curious human and an artificial intelligence assistant. "203 "The assistant gives helpful, detailed, and polite answers to the human's questions.",204 roles=("Human", "Assistant"),205 messages=(206 ("Human", "What are the key differences between renewable and non-renewable energy sources?"),207 ("Assistant",208 "Renewable energy sources are those that can be replenished naturally in a relatively "209 "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "210 "Non-renewable energy sources, on the other hand, are finite and will eventually be "211 "depleted, such as coal, oil, and natural gas. Here are some key differences between "212 "renewable and non-renewable energy sources:\n"213 "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "214 "energy sources are finite and will eventually run out.\n"215 "2. Environmental impact: Renewable energy sources have a much lower environmental impact "216 "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "217 "and other negative effects.\n"218 "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "219 "have lower operational costs than non-renewable sources.\n"220 "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "221 "locations than non-renewable sources.\n"222 "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "223 "situations and needs, while non-renewable sources are more rigid and inflexible.\n"224 "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "225 "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")226 ),227 offset=2,228 sep_style=SeparatorStyle.SINGLE,229 sep="###",230)231 232conv_vicuna_v1_1 = Conversation(233 system="A chat between a curious user and an artificial intelligence assistant. "234 "The assistant gives helpful, detailed, and polite answers to the user's questions.",235 roles=("USER", "ASSISTANT"),236 version="v1",237 messages=(),238 offset=0,239 sep_style=SeparatorStyle.TWO,240 sep=" ",241 sep2="</s>",242)243 244conv_mpt = Conversation(245 system="""<|im_start|>system246- You are a helpful language and vision assistant.247- You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.248- You should follow the instructions carefully and explain your answers in detail.""",249 roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),250 version="mpt",251 messages=(),252 offset=0,253 sep_style=SeparatorStyle.MPT,254 sep="<|im_end|>",255)256 257conv_mpt_text = Conversation(258 system="""<|im_start|>system259- You are a helpful assistant chatbot trained by MosaicML.260- You answer questions.261- You are excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.262- You are more than just an information source, you are also able to write poetry, short stories, and make jokes.""",263 roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),264 version="mpt",265 messages=(),266 offset=0,267 sep_style=SeparatorStyle.MPT,268 sep="<|im_end|>",269)270 271conv_bair_v1 = Conversation(272 system="BEGINNING OF CONVERSATION:",273 roles=("USER", "GPT"),274 messages=(),275 offset=0,276 sep_style=SeparatorStyle.TWO,277 sep=" ",278 sep2="</s>",279)280 281simple_conv = Conversation(282 system="A chat between a curious human and an artificial intelligence assistant. "283 "The assistant gives helpful, detailed, and polite answers to the human's questions.",284 roles=("Human", "Assistant"),285 messages=(286 ("Human", "Hi!"),287 ("Assistant", "Hi there! How can I help you today?")288 ),289 offset=2,290 sep_style=SeparatorStyle.SINGLE,291 sep="###",292)293 294simple_conv_multimodal = Conversation(295 system="You are LLaVA, a large language and vision assistant trained by UW Madison WAIV Lab."296 "You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."297 "Follow the instructions carefully and explain your answers in detail.",298 roles=("Human", "Assistant"),299 messages=(300 ("Human", "Hi!"),301 ("Assistant", "Hi there! How can I help you today?\n")302 ),303 offset=2,304 sep_style=SeparatorStyle.SINGLE,305 sep="###",306)307 308simple_conv_mpt_multimodal = Conversation(309 system="""<|im_start|>system310- You are LLaVA, a large language and vision assistant trained by UW Madison WAIV Lab.311- You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.312- You should follow the instructions carefully and explain your answers in detail.""",313 roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),314 version="mpt",315 messages=(),316 offset=0,317 sep_style=SeparatorStyle.MPT,318 sep="<|im_end|>",319)320 321simple_conv_legacy = Conversation(322 system="You are LLaVA, a large language model trained by UW Madison WAIV Lab."323 "You are designed to assist human with a variety of tasks using natural language."324 "Follow the instructions carefully.",325 roles=("Human", "Assistant"),326 messages=(327 ("Human", "Hi!\n\n### Response:"),328 ("Assistant", "Hi there! How can I help you today?\n")329 ),330 offset=2,331 sep_style=SeparatorStyle.SINGLE,332 sep="###",333)334 335conv_llava_v1 = Conversation(336 system="You are LLaVA, a large language and vision assistant trained by UW Madison WAIV Lab."337 "You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."338 "Follow the instructions carefully and explain your answers in detail.",339 roles=("USER", "ASSISTANT"),340 version="v1",341 messages=(),342 offset=0,343 sep_style=SeparatorStyle.TWO,344 sep=" ",345 sep2="</s>",346)347 348default_conversation = conv_v1_2349conv_templates = {350 "default": conv_v1_2,351 "simple": simple_conv,352 "simple_legacy": simple_conv_legacy,353 "multimodal": simple_conv_multimodal,354 "mpt_multimodal": simple_conv_mpt_multimodal,355 "llava_v1": conv_llava_v1,356 357 # fastchat358 "v1": conv_v1_2,359 "bair_v1": conv_bair_v1,360 "vicuna_v1_1": conv_vicuna_v1_1,361 "mpt": conv_mpt,362 "mpt_text": conv_mpt_text,363}364 365 366if __name__ == "__main__":367 print(default_conversation.get_prompt())