AsirAsir commited on
Commit
5683233
·
verified ·
1 Parent(s): 51df7bc

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +181 -59
app.py CHANGED
@@ -1,63 +1,185 @@
 
 
 
 
1
  import gradio as gr
2
- from huggingface_hub import InferenceClient
3
-
4
- """
5
- For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
6
- """
7
- client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
8
-
9
-
10
- def respond(
11
- message,
12
- history: list[tuple[str, str]],
13
- system_message,
14
- max_tokens,
15
- temperature,
16
- top_p,
17
- ):
18
- messages = [{"role": "system", "content": system_message}]
19
-
20
- for val in history:
21
- if val[0]:
22
- messages.append({"role": "user", "content": val[0]})
23
- if val[1]:
24
- messages.append({"role": "assistant", "content": val[1]})
25
-
26
- messages.append({"role": "user", "content": message})
27
-
28
- response = ""
29
-
30
- for message in client.chat_completion(
31
- messages,
32
- max_tokens=max_tokens,
33
- stream=True,
34
- temperature=temperature,
35
- top_p=top_p,
36
- ):
37
- token = message.choices[0].delta.content
38
-
39
- response += token
40
- yield response
41
-
42
- """
43
- For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
44
- """
45
- demo = gr.ChatInterface(
46
- respond,
47
- additional_inputs=[
48
- gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
49
- gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
50
- gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
51
- gr.Slider(
52
- minimum=0.1,
53
- maximum=1.0,
54
- value=0.95,
55
- step=0.05,
56
- label="Top-p (nucleus sampling)",
57
- ),
58
- ],
59
  )
 
 
60
 
 
61
 
62
- if __name__ == "__main__":
63
- demo.launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict
2
+ from typing import List
3
+ from typing import Tuple
4
+ from typing import Union
5
  import gradio as gr
6
+ import torch
7
+ import argparse
8
+ from threading import Thread
9
+ from transformers import (
10
+ AutoModelForCausalLM,
11
+ AutoTokenizer,
12
+ TextIteratorStreamer,
13
+ GenerationConfig
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  )
15
+ import warnings
16
+ import spaces
17
 
18
+ warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
19
 
20
+ ModelType = Union[PreTrainedModel]
21
+ TokenizerType = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
22
+
23
+ MODEL_PATH = os.environ.get('MODEL_PATH', 'IndexTeam/Index-1.9B-Chat')
24
+ TOKENIZER_PATH = os.environ.get("TOKENIZER_PATH", MODEL_PATH)
25
+
26
+ tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH, trust_remote_code=True)
27
+ model = AutoModelForCausalLM.from_pretrained(MODEL_PATH,
28
+ torch_dtype=torch.bfloat16,
29
+ device_map="auto",
30
+ trust_remote_code=True)
31
+
32
+ def _resolve_path(path: Union[str, Path]) -> Path:
33
+ return Path(path).expanduser().resolve()
34
+
35
+ @spaces.GPU
36
+ def hf_gen(dialog: List, top_k, top_p, temperature, repetition_penalty, max_dec_len):
37
+ """generate model output with huggingface api
38
+
39
+ Args:
40
+ query (str): actual model input.
41
+ top_p (float): only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation.
42
+ temperature (float): Strictly positive float value used to modulate the logits distribution.
43
+ max_dec_len (int): The maximum numbers of tokens to generate.
44
+
45
+ Yields:
46
+ str: real-time generation results of hf model
47
+ """
48
+ inputs = tokenizer.apply_chat_template(dialog, tokenize=False, add_generation_prompt=False)
49
+ enc = tokenizer(inputs, return_tensors="pt").to("cuda")
50
+ streamer = TextIteratorStreamer(tokenizer, **tokenizer.init_kwargs)
51
+ generation_kwargs = dict(
52
+ enc,
53
+ do_sample=True,
54
+ top_k=int(top_k),
55
+ top_p=float(top_p),
56
+ temperature=float(temperature),
57
+ repetition_penalty=float(repetition_penalty),
58
+ max_new_tokens=int(max_dec_len),
59
+ pad_token_id=tokenizer.eos_token_id,
60
+ streamer=streamer,
61
+ )
62
+ thread = Thread(target=model.generate, kwargs=generation_kwargs)
63
+ thread.start()
64
+ answer = ""
65
+ for new_text in streamer:
66
+ answer += new_text
67
+ yield answer[len(inputs):]
68
+
69
+ @spaces.GPU
70
+ def generate(chat_history: List, query, top_k, top_p, temperature, repetition_penalty, max_dec_len, system_message):
71
+ """generate after hitting "submit" button
72
+
73
+ Args:
74
+ chat_history (List): [[q_1, a_1], [q_2, a_2], ..., [q_n, a_n]]. list that stores all QA records
75
+ query (str): query of current round
76
+ top_p (float): only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation.
77
+ temperature (float): strictly positive float value used to modulate the logits distribution.
78
+ max_dec_len (int): The maximum numbers of tokens to generate.
79
+
80
+ Yields:
81
+ List: [[q_1, a_1], [q_2, a_2], ..., [q_n, a_n], [q_n+1, a_n+1]]. chat_history + QA of current round.
82
+ """
83
+ assert query != "", "Input must not be empty!!!"
84
+ # apply chat template
85
+ model_input = []
86
+ if system_message:
87
+ model_input.append({
88
+ "role": "system",
89
+ "content": system_message
90
+ })
91
+ for q, a in chat_history:
92
+ model_input.append({"role": "user", "content": q})
93
+ model_input.append({"role": "assistant", "content": a})
94
+ model_input.append({"role": "user", "content": query})
95
+ # yield model generation
96
+ chat_history.append([query, ""])
97
+ for answer in hf_gen(model_input, top_k, top_p, temperature, repetition_penalty, max_dec_len):
98
+ # chat_history[-1][1] = answer.strip("</s>")
99
+ chat_history[-1][1] = answer.strip(tokenizer.eos_token)
100
+ yield gr.update(value=""), chat_history
101
+
102
+ @spaces.GPU
103
+ def regenerate(chat_history: List, top_k, top_p, temperature, repetition_penalty, max_dec_len, system_message):
104
+ """re-generate the answer of last round's query
105
+
106
+ Args:
107
+ chat_history (List): [[q_1, a_1], [q_2, a_2], ..., [q_n, a_n]]. list that stores all QA records
108
+ top_p (float): only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation.
109
+ temperature (float): strictly positive float value used to modulate the logits distribution.
110
+ max_dec_len (int): The maximum numbers of tokens to generate.
111
+
112
+ Yields:
113
+ List: [[q_1, a_1], [q_2, a_2], ..., [q_n, a_n]]. chat_history
114
+ """
115
+ assert len(chat_history) >= 1, "History is empty. Nothing to regenerate!!"
116
+ # apply chat template
117
+ model_input = []
118
+ if system_message:
119
+ model_input.append({
120
+ "role": "system",
121
+ "content": system_message
122
+ })
123
+ for q, a in chat_history[:-1]:
124
+ model_input.append({"role": "user", "content": q})
125
+ model_input.append({"role": "assistant", "content": a})
126
+ model_input.append({"role": "user", "content": chat_history[-1][0]})
127
+ # yield model generation
128
+ for answer in hf_gen(model_input, top_k, top_p, temperature, repetition_penalty, max_dec_len):
129
+ # chat_history[-1][1] = answer.strip("</s>")
130
+ chat_history[-1][1] = answer.strip(tokenizer.eos_token)
131
+ yield gr.update(value=""), chat_history
132
+
133
+
134
+ def clear_history():
135
+ """clear all chat history
136
+
137
+ Returns:
138
+ List: empty chat history
139
+ """
140
+ torch.cuda.empty_cache()
141
+ return []
142
+
143
+
144
+ def reverse_last_round(chat_history):
145
+ """reverse last round QA and keep the chat history before
146
+
147
+ Args:
148
+ chat_history (List): [[q_1, a_1], [q_2, a_2], ..., [q_n, a_n]]. list that stores all QA records
149
+
150
+ Returns:
151
+ List: [[q_1, a_1], [q_2, a_2], ..., [q_n-1, a_n-1]]. chat_history without last round.
152
+ """
153
+ assert len(chat_history) >= 1, "History is empty. Nothing to reverse!!"
154
+ return chat_history[:-1]
155
+
156
+ # launch gradio demo
157
+ with gr.Blocks(theme="soft") as demo:
158
+ gr.Markdown("""# Index-1.9B Gradio Demo""")
159
+
160
+ with gr.Row():
161
+ with gr.Column(scale=1):
162
+ top_k = gr.Slider(1, 10, value=5, step=1, label="top_k")
163
+ top_p = gr.Slider(0, 1, value=0.8, step=0.1, label="top_p")
164
+ temperature = gr.Slider(0.1, 2.0, value=0.3, step=0.1, label="temperature")
165
+ repetition_penalty = gr.Slider(0.1, 2.0, value=1.1, step=0.1, label="repetition_penalty")
166
+ max_dec_len = gr.Slider(1, 4096, value=1024, step=1, label="max_dec_len")
167
+ with gr.Row():
168
+ system_message = gr.Textbox(label="System Message", placeholder="Input your system message", value="你是由哔哩哔哩自主研发的大语言模型,名为“Index”。你能够根据用户传入的信息,帮助用户完成指定的任务,并生成恰当的、符合要求的回复。")
169
+ with gr.Column(scale=10):
170
+ chatbot = gr.Chatbot(bubble_full_width=False, height=500, label='Index-1.9B')
171
+ user_input = gr.Textbox(label="User", placeholder="Input your query here!", lines=8)
172
+ with gr.Row():
173
+ submit = gr.Button("🚀 Submit")
174
+ clear = gr.Button("🧹 Clear")
175
+ regen = gr.Button("🔄 Regenerate")
176
+ reverse = gr.Button("⬅️ Reverse")
177
+
178
+ submit.click(generate, inputs=[chatbot, user_input, top_k, top_p, temperature, repetition_penalty, max_dec_len, system_message],
179
+ outputs=[user_input, chatbot])
180
+ regen.click(regenerate, inputs=[chatbot, top_k, top_p, temperature, repetition_penalty, max_dec_len, system_message],
181
+ outputs=[user_input, chatbot])
182
+ clear.click(clear_history, inputs=[], outputs=[chatbot])
183
+ reverse.click(reverse_last_round, inputs=[chatbot], outputs=[chatbot])
184
+
185
+ demo.queue().launch()