yuhuili commited on
Commit
06c0ba9
1 Parent(s): f077843

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +262 -0
app.py ADDED
@@ -0,0 +1,262 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import time
3
+
4
+ #os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
5
+ import gradio as gr
6
+ import argparse
7
+ from model.ex_model import ExModel
8
+ import torch
9
+ from fastchat.model import get_conversation_template
10
+ import re
11
+
12
+
13
+ def truncate_list(lst, num):
14
+ if num not in lst:
15
+ return lst
16
+ first_index = lst.index(num)
17
+ return lst[:first_index + 1]
18
+
19
+
20
+
21
+
22
+
23
+ def find_list_markers(text):
24
+ pattern = re.compile(r'(?m)(^\d+\.\s|\n)')
25
+ matches = pattern.finditer(text)
26
+ return [(match.start(), match.end()) for match in matches]
27
+
28
+
29
+ def checkin(pointer,start,marker):
30
+ for b,e in marker:
31
+ if b<=pointer<e:
32
+ return True
33
+ if b<=start<e:
34
+ return True
35
+ return False
36
+
37
+ def highlight_text(text, text_list,color="black"):
38
+ pointer = 0
39
+ result = ""
40
+ markers=find_list_markers(text)
41
+
42
+ for sub_text in text_list:
43
+
44
+ start = text.find(sub_text, pointer)
45
+ if start==-1:
46
+ continue
47
+ end = start + len(sub_text)
48
+
49
+ if checkin(pointer,start,markers):
50
+ result += text[pointer:start]
51
+ else:
52
+ result += f"<span style='color: {color};'>{text[pointer:start]}</span>"
53
+
54
+
55
+ result += sub_text
56
+
57
+
58
+ pointer = end
59
+
60
+ if pointer < len(text):
61
+ result += f"<span style='color: {color};'>{text[pointer:]}</span>"
62
+
63
+ return result
64
+
65
+
66
+ def warmup(model):
67
+ conv = get_conversation_template(args.model_type)
68
+
69
+ if args.model_type == "llama-2-chat":
70
+ sys_p = "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.\n\nIf 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."
71
+ conv.system_message = sys_p
72
+ conv.append_message(conv.roles[0], "Hello")
73
+ conv.append_message(conv.roles[1], None)
74
+ prompt = conv.get_prompt()
75
+ if args.model_type == "llama-2-chat":
76
+ prompt += " "
77
+ input_ids = model.tokenizer([prompt]).input_ids
78
+ input_ids = torch.as_tensor(input_ids).cuda()
79
+ for output_ids in model.ex_generate(input_ids):
80
+ ol=output_ids.shape[1]
81
+
82
+ def bot(history, session_state):
83
+ temperature = 0.5
84
+ top_p = 0.9
85
+ if not history:
86
+ return history,"0.00 tokens/s","0.00",session_state
87
+ pure_history=session_state.get("pure_history",[])
88
+ assert args.model_type == "llama-2-chat" or "vicuna"
89
+ conv = get_conversation_template(args.model_type)
90
+
91
+ if args.model_type == "llama-2-chat":
92
+ sys_p = "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.\n\nIf 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."
93
+ conv.system_message = sys_p
94
+
95
+ for query, response in pure_history:
96
+ conv.append_message(conv.roles[0], query)
97
+ if args.model_type == "llama-2-chat" and response:
98
+ response = " " + response
99
+ conv.append_message(conv.roles[1], response)
100
+
101
+ prompt = conv.get_prompt()
102
+
103
+ if args.model_type == "llama-2-chat":
104
+ prompt += " "
105
+
106
+ input_ids = model.tokenizer([prompt]).input_ids
107
+ input_ids = torch.as_tensor(input_ids).cuda()
108
+ input_len = input_ids.shape[1]
109
+ naive_text = []
110
+ cu_len = input_len
111
+ totaltime=0
112
+ start_time=time.time()
113
+ total_ids=0
114
+
115
+
116
+ for output_ids in model.ex_generate(input_ids, temperature=temperature, top_p=top_p,
117
+ max_steps=args.max_new_token):
118
+ totaltime+=(time.time()-start_time)
119
+ total_ids+=1
120
+ decode_ids = output_ids[0, input_len:].tolist()
121
+ decode_ids = truncate_list(decode_ids, model.tokenizer.eos_token_id)
122
+ text = model.tokenizer.decode(decode_ids, skip_special_tokens=True, spaces_between_special_tokens=False,
123
+ clean_up_tokenization_spaces=True, )
124
+ naive_text.append(model.tokenizer.decode(output_ids[0, cu_len], skip_special_tokens=True,
125
+ spaces_between_special_tokens=False,
126
+ clean_up_tokenization_spaces=True, ))
127
+
128
+ cu_len = output_ids.shape[1]
129
+ colored_text = highlight_text(text, naive_text, "orange")
130
+ #if highlight_ExInfer:
131
+ history[-1][1] = colored_text
132
+ # else:
133
+ # history[-1][1] = text
134
+ pure_history[-1][1] = text
135
+ session_state["pure_history"]=pure_history
136
+ new_tokens = cu_len-input_len
137
+ yield history,f"{new_tokens/totaltime:.2f} tokens/s",f"{new_tokens/total_ids:.2f}",session_state
138
+ start_time = time.time()
139
+
140
+
141
+
142
+
143
+
144
+ def user(user_message, history,session_state):
145
+ if history==None:
146
+ history=[]
147
+ pure_history = session_state.get("pure_history", [])
148
+ pure_history += [[user_message, None]]
149
+ session_state["pure_history"] = pure_history
150
+ return "", history + [[user_message, None]],session_state
151
+
152
+
153
+ def regenerate(history,session_state):
154
+
155
+ try:
156
+
157
+ if not history:
158
+ return history, None,"0.00 tokens/s","0.00",session_state
159
+ pure_history = session_state.get("pure_history", [])
160
+ try:
161
+ pure_history[-1][-1] = None
162
+ except:
163
+ print(1)
164
+ session_state["pure_history"]=pure_history
165
+ if len(history) > 1: # Check if there's more than one entry in history (i.e., at least one bot response)
166
+ new_history = history[:-1] # Remove the last bot response
167
+ last_user_message = history[-1][0] # Get the last user message
168
+ return new_history + [[last_user_message, None]], None,"0.00 tokens/s","0.00",session_state
169
+ history[-1][1] = None
170
+ return history, None,"0.00 tokens/s","0.00",session_state
171
+
172
+ except:
173
+ print(2)
174
+ return history, None, "0.00 tokens/s", "0.00", session_state
175
+
176
+
177
+ def clear(history,session_state):
178
+ pure_history = session_state.get("pure_history", [])
179
+ pure_history = []
180
+ session_state["pure_history"] = pure_history
181
+ return [],"0.00 tokens/s","0.00",session_state
182
+
183
+
184
+
185
+
186
+ parser = argparse.ArgumentParser()
187
+ parser.add_argument(
188
+ "--ex-model-path",
189
+ type=str,
190
+ default="lmsys/vicuna-7b-v1.3",
191
+ help="The path to the weights. This can be a local folder or a Hugging Face repo ID.",
192
+ )
193
+ parser.add_argument("--base-model-path", type=str, default="/home/lyh/weights/hf/vicuna_v13/13B/",
194
+ help="path of basemodel, huggingface project or local path")
195
+ parser.add_argument(
196
+ "--load-in-8bit", action="store_true", help="Use 8-bit quantization"
197
+ )
198
+ parser.add_argument(
199
+ "--load-in-4bit", action="store_true", help="Use 4-bit quantization"
200
+ )
201
+ parser.add_argument("--model-type", type=str, default="vicuna", help="llama-2-chat or vicuna, for chat template")
202
+ parser.add_argument(
203
+ "--max-new-token",
204
+ type=int,
205
+ default=512,
206
+ help="The maximum number of new generated tokens.",
207
+ )
208
+ args = parser.parse_args()
209
+
210
+ model = ExModel.from_pretrained(
211
+ base_model_path=args.base_model_path,
212
+ ex_model_path=args.ex_model_path,
213
+ torch_dtype=torch.float16,
214
+ low_cpu_mem_usage=True,
215
+ load_in_4bit=args.load_in_4bit,
216
+ load_in_8bit=args.load_in_8bit,
217
+ device_map="auto"
218
+ )
219
+ model.eval()
220
+ warmup(model)
221
+
222
+ custom_css = """
223
+ #speed textarea {
224
+ color: red;
225
+ font-size: 30px;
226
+ }"""
227
+
228
+
229
+ with gr.Blocks(css=custom_css) as demo:
230
+ gs=gr.State({"pure_history":[]})
231
+ gr.Markdown('''## EAGLE Chatbot''')
232
+ with gr.Row():
233
+ speed_box = gr.Textbox(label="Speed", elem_id="speed", interactive=False, value="0.00 tokens/s")
234
+ compression_box = gr.Textbox(label="Compression Ratio", elem_id="speed", interactive=False, value="0.00")
235
+ note1 = gr.Markdown(show_label=False, interactive=False,
236
+ value='''The Compression Ratio is defined as the number of generated tokens divided by the number of forward passes in the original LLM. The original LLM is Vicuna 13B, with inference conducted on RTX 3090 GPUs and at a precision of fp16.''')
237
+ note=gr.Markdown(show_label=False,interactive=False,value='''The tokens that EAGLE correctly guesses will be highlighted in orange. Note: This highlighting may lead to special formatting rendering issues in some instances, particularly when generating code.''')
238
+
239
+
240
+ chatbot = gr.Chatbot(height=600,show_label=False)
241
+
242
+
243
+ msg = gr.Textbox(label="Your input")
244
+ with gr.Row():
245
+ send_button = gr.Button("Send")
246
+ stop_button = gr.Button("Stop")
247
+ regenerate_button = gr.Button("Regenerate")
248
+ clear_button = gr.Button("Clear")
249
+ enter_event=msg.submit(user, [msg, chatbot,gs], [msg, chatbot,gs], queue=True).then(
250
+ bot, [chatbot,gs ], [chatbot,speed_box,compression_box,gs]
251
+ )
252
+ clear_button.click(clear, [chatbot,gs], [chatbot,speed_box,compression_box,gs], queue=True)
253
+
254
+ send_event=send_button.click(user, [msg, chatbot,gs], [msg, chatbot,gs],queue=True).then(
255
+ bot, [chatbot,gs ], [chatbot,speed_box,compression_box,gs]
256
+ )
257
+ regenerate_event=regenerate_button.click(regenerate, [chatbot,gs], [chatbot, msg,speed_box,compression_box,gs],queue=True).then(
258
+ bot, [chatbot,gs ], [chatbot,speed_box,compression_box,gs]
259
+ )
260
+ stop_button.click(fn=None, inputs=None, outputs=None, cancels=[send_event,regenerate_event,enter_event])
261
+ demo.queue()
262
+ demo.launch(share=True)