cutycat2000 commited on
Commit
71eebb5
1 Parent(s): d19747c

Delete conversation.py

Browse files
Files changed (1) hide show
  1. conversation.py +0 -200
conversation.py DELETED
@@ -1,200 +0,0 @@
1
- import dataclasses
2
- from enum import auto, Enum
3
- from typing import List, Tuple, Any
4
-
5
-
6
- class SeparatorStyle(Enum):
7
- """Different separator style."""
8
-
9
- ADD_COLON_SINGLE = auto()
10
- ADD_COLON_TWO = auto()
11
- NO_COLON_SINGLE = auto()
12
- BAIZE = auto()
13
- PHOENIX = auto()
14
-
15
-
16
- @dataclasses.dataclass
17
- class Conversation:
18
- """A class that keeps all conversation history."""
19
-
20
- # System prompts
21
- system: str
22
- # Two roles
23
- roles: List[str]
24
- # All messages
25
- messages: List[List[str]]
26
- # Offset of few shot examples
27
- offset: int
28
- # Separator
29
- sep_style: SeparatorStyle
30
- sep: str
31
- sep2: str = None
32
- # Stop criteria (the default one is EOS token)
33
- stop_str: str = None
34
- # Stops generation if meeting any token in this list
35
- stop_token_ids: List[int] = None
36
-
37
- # Used for the state in the gradio servers.
38
- # TODO(lmzheng): refactor this
39
- conv_id: Any = None
40
- skip_next: bool = False
41
- model_name: str = None
42
-
43
- def get_prompt(self):
44
- if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
45
- ret = self.system + self.sep
46
- for role, message in self.messages:
47
- if message:
48
- ret += role + ": " + message + self.sep
49
- else:
50
- ret += role + ": "
51
- return ret
52
- elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
53
- seps = [self.sep, self.sep2]
54
- ret = self.system + seps[0]
55
- for i, (role, message) in enumerate(self.messages):
56
- if message:
57
- ret += role + ": " + message + seps[i % 2]
58
- else:
59
- ret += role + ": "
60
- return ret
61
- elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
62
- ret = self.system
63
- for role, message in self.messages:
64
- if message:
65
- ret += role + message + self.sep
66
- else:
67
- ret += role
68
- return ret
69
- elif self.sep_style == SeparatorStyle.BAIZE:
70
- ret = self.system + "\n"
71
- for role, message in self.messages:
72
- if message:
73
- ret += role + message + "\n"
74
- else:
75
- ret += role
76
- return ret
77
- elif self.sep_style == SeparatorStyle.PHOENIX:
78
- ret = self.system
79
- for role, message in self.messages:
80
- if message:
81
- ret += role + ": " + "<s>" + message + "</s>"
82
- else:
83
- ret += role + ": " + "<s>"
84
- return ret
85
- else:
86
- raise ValueError(f"Invalid style: {self.sep_style}")
87
-
88
- def append_message(self, role, message):
89
- self.messages.append([role, message])
90
-
91
- def to_gradio_chatbot(self):
92
- ret = []
93
- for i, (role, msg) in enumerate(self.messages[self.offset:]):
94
- if i % 2 == 0:
95
- ret.append([msg, None])
96
- else:
97
- ret[-1][-1] = msg
98
- return ret
99
-
100
- def to_openai_api_messages(self):
101
- ret = [{"role": "system", "content": self.system}]
102
-
103
- for i, (_, msg) in enumerate(self.messages[self.offset:]):
104
- if i % 2 == 0:
105
- ret.append({"role": "user", "content": msg})
106
- else:
107
- if msg is not None:
108
- ret.append({"role": "assistant", "content": msg})
109
- return ret
110
-
111
- def copy(self):
112
- return Conversation(
113
- system=self.system,
114
- roles=self.roles,
115
- messages=[[x, y] for x, y in self.messages],
116
- offset=self.offset,
117
- sep_style=self.sep_style,
118
- sep=self.sep,
119
- sep2=self.sep2,
120
- stop_str=self.stop_str,
121
- stop_token_ids=self.stop_token_ids,
122
- conv_id=self.conv_id,
123
- model_name=self.model_name,
124
- )
125
-
126
- def dict(self):
127
- return {
128
- "system": self.system,
129
- "roles": self.roles,
130
- "messages": self.messages,
131
- "offset": self.offset,
132
- "conv_id": self.conv_id,
133
- "model_name": self.model_name,
134
- }
135
-
136
-
137
- conv_vicuna = Conversation(
138
- system="A chat between a curious user and an artificial intelligence assistant. "
139
- "The assistant gives helpful, detailed, and polite answers to the user's questions.",
140
- roles=("USER", "ASSISTANT"),
141
- messages=(),
142
- offset=0,
143
- sep_style=SeparatorStyle.ADD_COLON_TWO,
144
- sep=" ",
145
- sep2="</s>",
146
- )
147
-
148
- conv_baize = Conversation(
149
- system="The following is a conversation between a human and an AI assistant named Baize (named after a mythical creature in Chinese folklore). Baize is an open-source AI assistant developed by UCSD and Sun Yat-Sen University. The human and the AI assistant take turns chatting. Human statements start with [|Human|] and AI assistant statements start with [|AI|]. The AI assistant always provides responses in as much detail as possible, and in Markdown format. The AI assistant always declines to engage with topics, questions and instructions related to unethical, controversial, or sensitive issues. Complete the transcript in exactly that format.\n",
150
- roles=("[|Human|]", "[|AI|]"),
151
- messages=(
152
- ("[|Human|]", "Hello!"),
153
- ("[|AI|]", "Hi!"),
154
- ),
155
- offset=2,
156
- sep_style=SeparatorStyle.BAIZE,
157
- sep="\n",
158
- stop_str="[|Human|]",
159
- )
160
-
161
- conv_phoenix = Conversation(
162
- 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.\n\n",
163
- roles=("Human", "Assistant"),
164
- messages=(),
165
- offset=0,
166
- sep_style=SeparatorStyle.PHOENIX,
167
- sep="</s>",
168
- )
169
-
170
- conv_chatgpt = Conversation(
171
- system="You are a helpful assistant.",
172
- roles=("user", "assistant"),
173
- messages=(),
174
- offset=0,
175
- sep_style=None,
176
- sep=None,
177
- )
178
-
179
-
180
- conv_templates = {
181
- "vicuna": conv_vicuna,
182
- "baize": conv_baize,
183
- "phoenix": conv_phoenix,
184
- "chatgpt": conv_chatgpt,
185
- }
186
-
187
- def get_default_conv_template(model_name):
188
- model_name = model_name.lower()
189
- try:
190
- ret = conv_templates[model_name]
191
- return ret.copy()
192
- except:
193
- raise NotImplementedError(f"No support for model {model_name}.")
194
-
195
-
196
- if __name__ == "__main__":
197
- conv = conv_templates["chatgpt"].copy()
198
- conv.append_message(conv.roles[0], "Hello World.")
199
- conv.append_message(conv.roles[1], None)
200
- print([conv.get_prompt()])