simplecloud/VidChain-exercise

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    PLAIN = auto()12    LLAMA_2 = auto()13 14 15@dataclasses.dataclass16class Conversation:17    """A class that keeps all conversation history."""18    system: str19    roles: List[str]20    messages: List[List[str]]21    offset: int22    sep_style: SeparatorStyle = SeparatorStyle.SINGLE23    sep: str = "###"24    sep2: str = None25    version: str = "Unknown"26 27    skip_next: bool = False28 29    def get_prompt(self):30        messages = self.messages31        if len(messages) > 0 and type(messages[0][1]) is tuple:32            messages = self.messages.copy()33            init_role, init_msg = messages[0].copy()34            init_msg = init_msg[0].replace("<image>", "").strip()35            if 'mmtag' in self.version:36                messages[0] = (init_role, init_msg)37                messages.insert(0, (self.roles[0], "<Image><image></Image>"))38                messages.insert(1, (self.roles[1], "Received."))39            else:40                messages[0] = (init_role, "<image>\n" + init_msg)41 42        if self.sep_style == SeparatorStyle.SINGLE:43            ret = self.system + self.sep44            for role, message in messages:45                if message:46                    if type(message) is tuple:47                        message, _, _ = message48                    ret += role + ": " + message + self.sep49                else:50                    ret += role + ":"51        elif self.sep_style == SeparatorStyle.TWO:52            seps = [self.sep, self.sep2]53            ret = self.system + seps[0]54            for i, (role, message) in enumerate(messages):55                if message:56                    if type(message) is tuple:57                        message, _, _ = message58                    ret += role + ": " + message + seps[i % 2]59                else:60                    ret += role + ":"61        elif self.sep_style == SeparatorStyle.MPT:62            ret = self.system + self.sep63            for role, message in messages:64                if message:65                    if type(message) is tuple:66                        message, _, _ = message67                    ret += role + message + self.sep68                else:69                    ret += role70        elif self.sep_style == SeparatorStyle.LLAMA_2:71            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n"72            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"73            ret = ""74 75            for i, (role, message) in enumerate(messages):76                if i == 0:77                    assert message, "first message should not be none"78                    assert role == self.roles[0], "first message should come from user"79                if message:80                    if type(message) is tuple:81                        message, _, _ = message82                    if i == 0: message = wrap_sys(self.system) + message83                    if i % 2 == 0:84                        message = wrap_inst(message)85                        ret += self.sep + message86                    else:87                        ret += " " + message + " " + self.sep288                else:89                    ret += ""90            ret = ret.lstrip(self.sep)91        elif self.sep_style == SeparatorStyle.PLAIN:92            seps = [self.sep, self.sep2]93            ret = self.system94            for i, (role, message) in enumerate(messages):95                if message:96                    if type(message) is tuple:97                        message, _, _ = message98                    ret += message + seps[i % 2]99                else:100                    ret += ""101        else:102            raise ValueError(f"Invalid style: {self.sep_style}")103 104        return ret105 106    def append_message(self, role, message):107        self.messages.append([role, message])108 109    def remove_message(self, role, message):110        self.messages.remove([role, message])111 112    def get_images(self, return_pil=False):113        images = []114        for i, (role, msg) in enumerate(self.messages[self.offset:]):115            if i % 2 == 0:116                if type(msg) is tuple:117                    import base64118                    from io import BytesIO119                    from PIL import Image120                    msg, image, image_process_mode = msg121                    if image_process_mode == "Pad":122                        def expand2square(pil_img, background_color=(122, 116, 104)):123                            width, height = pil_img.size124                            if width == height:125                                return pil_img126                            elif width > height:127                                result = Image.new(pil_img.mode, (width, width), background_color)128                                result.paste(pil_img, (0, (width - height) // 2))129                                return result130                            else:131                                result = Image.new(pil_img.mode, (height, height), background_color)132                                result.paste(pil_img, ((height - width) // 2, 0))133                                return result134                        image = expand2square(image)135                    elif image_process_mode == "Crop":136                        pass137                    elif image_process_mode == "Resize":138                        image = image.resize((336, 336))139                    else:140                        raise ValueError(f"Invalid image_process_mode: {image_process_mode}")141                    max_hw, min_hw = max(image.size), min(image.size)142                    aspect_ratio = max_hw / min_hw143                    max_len, min_len = 800, 400144                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))145                    longest_edge = int(shortest_edge * aspect_ratio)146                    W, H = image.size147                    if H > W:148                        H, W = longest_edge, shortest_edge149                    else:150                        H, W = shortest_edge, longest_edge151                    image = image.resize((W, H))152                    if return_pil:153                        images.append(image)154                    else:155                        buffered = BytesIO()156                        image.save(buffered, format="PNG")157                        img_b64_str = base64.b64encode(buffered.getvalue()).decode()158                        images.append(img_b64_str)159        return images160 161    def to_gradio_chatbot(self):162        ret = []163        for i, (role, msg) in enumerate(self.messages[self.offset:]):164            if i % 2 == 0:165                if type(msg) is tuple:166                    import base64167                    from io import BytesIO168                    msg, image, image_process_mode = msg169                    max_hw, min_hw = max(image.size), min(image.size)170                    aspect_ratio = max_hw / min_hw171                    max_len, min_len = 800, 400172                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))173                    longest_edge = int(shortest_edge * aspect_ratio)174                    W, H = image.size175                    if H > W:176                        H, W = longest_edge, shortest_edge177                    else:178                        H, W = shortest_edge, longest_edge179                    image = image.resize((W, H))180                    buffered = BytesIO()181                    image.save(buffered, format="JPEG")182                    img_b64_str = base64.b64encode(buffered.getvalue()).decode()183                    img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'184                    ret.append([img_str, None])185                    msg = msg.replace('<image>', '').strip()186                    if len(msg) > 0:187                        ret.append([msg, None])188                else:189                    ret.append([msg, None])190            else:191                ret[-1][-1] = msg192        return ret193 194    def copy(self):195        return Conversation(196            system=self.system,197            roles=self.roles,198            messages=[[x, y] for x, y in self.messages],199            offset=self.offset,200            sep_style=self.sep_style,201            sep=self.sep,202            sep2=self.sep2,203            version=self.version)204 205    def dict(self):206        if len(self.get_images()) > 0:207            return {208                "system": self.system,209                "roles": self.roles,210                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],211                "offset": self.offset,212                "sep": self.sep,213                "sep2": self.sep2,214            }215        return {216            "system": self.system,217            "roles": self.roles,218            "messages": self.messages,219            "offset": self.offset,220            "sep": self.sep,221            "sep2": self.sep2,222        }223 224 225conv_vicuna_v0 = Conversation(226    system="A chat between a curious human and an artificial intelligence assistant. "227           "The assistant gives helpful, detailed, and polite answers to the human's questions.",228    roles=("Human", "Assistant"),229    messages=(230        ("Human", "What are the key differences between renewable and non-renewable energy sources?"),231        ("Assistant",232            "Renewable energy sources are those that can be replenished naturally in a relatively "233            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "234            "Non-renewable energy sources, on the other hand, are finite and will eventually be "235            "depleted, such as coal, oil, and natural gas. Here are some key differences between "236            "renewable and non-renewable energy sources:\n"237            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "238            "energy sources are finite and will eventually run out.\n"239            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "240            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "241            "and other negative effects.\n"242            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "243            "have lower operational costs than non-renewable sources.\n"244            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "245            "locations than non-renewable sources.\n"246            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "247            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"248            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "249            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")250    ),251    offset=2,252    sep_style=SeparatorStyle.SINGLE,253    sep="###",254)255 256conv_vicuna_v1 = Conversation(257    system="A chat between a curious user and an artificial intelligence assistant. "258    "The assistant gives helpful, detailed, and polite answers to the user's questions. Avoid repeating sentences or words.",259    roles=("USER", "ASSISTANT"),260    version="v1",261    messages=(),262    offset=0,263    sep_style=SeparatorStyle.TWO,264    sep=" ",265    sep2="</s>",266)267 268conv_llama_2 = Conversation(269    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.270 271If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",272    roles=("USER", "ASSISTANT"),273    version="llama_v2",274    messages=(),275    offset=0,276    sep_style=SeparatorStyle.LLAMA_2,277    sep="<s>",278    sep2="</s>",279)280 281conv_llava_llama_2 = Conversation(282    system="You are a helpful language and vision assistant. "283           "You are able to understand the visual content that the user provides, "284           "and assist the user with a variety of tasks using natural language.",285    roles=("USER", "ASSISTANT"),286    version="llama_v2",287    messages=(),288    offset=0,289    sep_style=SeparatorStyle.LLAMA_2,290    sep="<s>",291    sep2="</s>",292)293 294conv_mpt = Conversation(295    system="""<|im_start|>system296A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",297    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),298    version="mpt",299    messages=(),300    offset=0,301    sep_style=SeparatorStyle.MPT,302    sep="<|im_end|>",303)304 305conv_llava_plain = Conversation(306    system="",307    roles=("", ""),308    messages=(309    ),310    offset=0,311    sep_style=SeparatorStyle.PLAIN,312    sep="\n",313)314 315conv_llava_v0 = Conversation(316    system="A chat between a curious human and an artificial intelligence assistant. "317           "The assistant gives helpful, detailed, and polite answers to the human's questions.",318    roles=("Human", "Assistant"),319    messages=(320        ("Human", "Hi!"),321        ("Assistant", "Hi there! How can I help you today?")322    ),323    offset=2,324    sep_style=SeparatorStyle.SINGLE,325    sep="###",326)327 328conv_llava_v0_mmtag = Conversation(329    system="A chat between a curious user and an artificial intelligence assistant. "330           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."331           "The visual content will be provided with the following format: <Image>visual content</Image>.",332    roles=("Human", "Assistant"),333    messages=(334    ),335    offset=0,336    sep_style=SeparatorStyle.SINGLE,337    sep="###",338    version="v0_mmtag",339)340 341conv_llava_v1 = Conversation(342    system="A chat between a curious human and an artificial intelligence assistant. "343           "The assistant gives helpful, detailed, and polite answers to the human's questions.",344    roles=("USER", "ASSISTANT"),345    version="v1",346    messages=(),347    offset=0,348    sep_style=SeparatorStyle.TWO,349    sep=" ",350    sep2="</s>",351)352 353conv_llava_v1_mmtag = Conversation(354    system="A chat between a curious user and an artificial intelligence assistant. "355           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."356           "The visual content will be provided with the following format: <Image>visual content</Image>.",357    roles=("USER", "ASSISTANT"),358    messages=(),359    offset=0,360    sep_style=SeparatorStyle.TWO,361    sep=" ",362    sep2="</s>",363    version="v1_mmtag",364)365 366default_conversation = conv_vicuna_v1367conv_templates = {368    "default": conv_vicuna_v0,369    "v0": conv_vicuna_v0,370    "v1": conv_vicuna_v1,371    "vicuna_v1": conv_vicuna_v1,372    "llama_2": conv_llama_2,373 374    "plain": conv_llava_plain,375    "v0_plain": conv_llava_plain,376    "llava_v0": conv_llava_v0,377    "v0_mmtag": conv_llava_v0_mmtag,378    "llava_v1": conv_llava_v1,379    "v1_mmtag": conv_llava_v1_mmtag,380    "llava_llama_2": conv_llava_llama_2,381 382    "mpt": conv_mpt,383}384 385 386if __name__ == "__main__":387    print(default_conversation.get_prompt())388