colossalai
89 строк · 2.5 Кб
1# Copyright 2023 lm-sys@FastChat
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import dataclasses16from enum import Enum, auto17from typing import List18
19
20class SeparatorStyle(Enum):21ADD_EOS_TOKEN = auto()22
23
24@dataclasses.dataclass25class Conversation:26system: str27roles: List[str]28messages: List[List[str]]29offset: int30sep_style: SeparatorStyle = SeparatorStyle.ADD_EOS_TOKEN31sep: str = "</s>"32
33skip_next: bool = False34
35def get_prompt(self):36if self.sep_style == SeparatorStyle.ADD_EOS_TOKEN:37ret = self.system38for role, message in self.messages:39if message:40ret += role + ": " + message + self.sep41else:42ret += role + ": "43return ret44else:45raise ValueError(f"Invalid style: {self.sep_style}")46
47def append_message(self, role, message):48self.messages.append([role, message])49
50def to_gradio_chatbot(self):51ret = []52for i, (role, msg) in enumerate(self.messages[self.offset :]):53if i % 2 == 0:54ret.append([msg, None])55else:56ret[-1][-1] = msg57return ret58
59def copy(self):60return Conversation(61system=self.system,62roles=self.roles,63messages=[[x, y] for x, y in self.messages],64offset=self.offset,65sep_style=self.sep_style,66sep=self.sep,67)68
69def dict(self):70return {71"system": self.system,72"roles": self.roles,73"messages": self.messages,74"offset": self.offset,75"sep": self.sep,76}77
78
79conv = Conversation(80system="A chat between a curious human and an artificial intelligence assistant. "81"The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",82roles=("Human", "Assistant"),83messages=(),84offset=0,85sep_style=SeparatorStyle.ADD_EOS_TOKEN,86sep="</s>",87)
88
89default_conversation = conv90