yilinw commited on
Commit
58023c8
1 Parent(s): f48d564

gradio_app file

Browse files
Files changed (1) hide show
  1. gradio_demo.py +474 -0
gradio_demo.py ADDED
@@ -0,0 +1,474 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import (
3
+ LlamaForCausalLM,
4
+ LlamaTokenizer,
5
+ StoppingCriteria,
6
+ )
7
+ import gradio as gr
8
+ import argparse
9
+ import os
10
+ from queue import Queue
11
+ from threading import Thread
12
+ import traceback
13
+ import gc
14
+
15
+
16
+ import torch
17
+ from auto_gptq import AutoGPTQForCausalLM
18
+ from langchain import HuggingFacePipeline, PromptTemplate
19
+ from langchain.chains import RetrievalQA
20
+ from langchain.document_loaders import PyPDFDirectoryLoader, DirectoryLoader
21
+ from langchain.embeddings import HuggingFaceInstructEmbeddings
22
+ from langchain.text_splitter import RecursiveCharacterTextSplitter
23
+ from langchain.vectorstores import Chroma
24
+ from pdf2image import convert_from_path
25
+ from transformers import AutoTokenizer, TextStreamer, pipeline
26
+
27
+ DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
28
+
29
+ # Parse command-line arguments
30
+ parser = argparse.ArgumentParser()
31
+ parser.add_argument(
32
+ '--base_model',
33
+ default=None,
34
+ type=str,
35
+ help='Base model path')
36
+ parser.add_argument('--lora_model', default=None, type=str,
37
+ help="If None, perform inference on the base model")
38
+ parser.add_argument(
39
+ '--tokenizer_path',
40
+ default=None,
41
+ type=str,
42
+ help='If None, lora model path or base model path will be used')
43
+ parser.add_argument(
44
+ '--gpus',
45
+ default="0",
46
+ type=str,
47
+ help='If None, cuda:0 will be used. Inference using multi-cards: --gpus=0,1,... ')
48
+ parser.add_argument('--share', default=True, help='Share gradio domain name')
49
+ parser.add_argument('--port', default=19324, type=int, help='Port of gradio demo')
50
+ parser.add_argument(
51
+ '--max_memory',
52
+ default=256,
53
+ type=int,
54
+ help='Maximum input prompt length, if exceeded model will receive prompt[-max_memory:]')
55
+ parser.add_argument(
56
+ '--load_in_8bit',
57
+ action='store_true',
58
+ help='Use 8 bit quantified model')
59
+ parser.add_argument(
60
+ '--only_cpu',
61
+ action='store_true',
62
+ help='Only use CPU for inference')
63
+ parser.add_argument(
64
+ '--alpha',
65
+ type=str,
66
+ default="1.0",
67
+ help="The scaling factor of NTK method, can be a float or 'auto'. ")
68
+ args = parser.parse_args()
69
+ if args.only_cpu is True:
70
+ args.gpus = ""
71
+
72
+ #from patches import apply_attention_patch, apply_ntk_scaling_patch
73
+ #apply_attention_patch(use_memory_efficient_attention=True)
74
+ #apply_ntk_scaling_patch(args.alpha)
75
+
76
+ # Set CUDA devices if available
77
+ os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
78
+
79
+
80
+ # Peft library can only import after setting CUDA devices
81
+ from peft import PeftModel
82
+
83
+
84
+ # Set up the required components: model and tokenizer
85
+
86
+ def setup():
87
+ global tokenizer, model, device, share, port, max_memory, vector_store
88
+ max_memory = args.max_memory
89
+ port = args.port
90
+ share = args.share
91
+ load_in_8bit = args.load_in_8bit
92
+ load_type = torch.float16
93
+ if torch.cuda.is_available():
94
+ device = torch.device(0)
95
+ else:
96
+ device = torch.device('cpu')
97
+ """
98
+ if args.tokenizer_path is None:
99
+ args.tokenizer_path = args.lora_model
100
+ if args.lora_model is None:
101
+ args.tokenizer_path = args.base_model
102
+ """
103
+
104
+
105
+ #先读取embedding模型
106
+ embeddings = HuggingFaceInstructEmbeddings(
107
+ model_name="BAAI/bge-large-en-v1.5", model_kwargs={"device": DEVICE}
108
+ )
109
+ #如果之前没有本地的faiss仓库,就把doc读取到向量库后,再把向量库保存到本地
110
+ if os.path.exists("/home/ywang/db")==False:
111
+ #=======加载知识库=======
112
+ loader = DirectoryLoader("kb")
113
+ docs = loader.load()
114
+ # splitting pdf into chunks with size of 1024
115
+ text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
116
+ texts = text_splitter.split_documents(docs)
117
+ vector_store = Chroma.from_documents(texts, embeddings, persist_directory="db")
118
+ #如果本地已经有faiss仓库了,说明之前已经保存过了,就直接读取
119
+ else:
120
+ vector_store=Chroma(persist_directory="db", embedding_function=embeddings)
121
+
122
+ model_name_or_path = "TheBloke/Llama-2-13B-chat-GPTQ"
123
+ model_basename = "model"
124
+
125
+ tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
126
+
127
+ base_model = AutoGPTQForCausalLM.from_quantized(
128
+ model_name_or_path,
129
+ revision="gptq-4bit-128g-actorder_True",
130
+ model_basename=model_basename,
131
+ use_safetensors=True,
132
+ trust_remote_code=True,
133
+ inject_fused_attention=False,
134
+ device=DEVICE,
135
+ quantize_config=None,
136
+ )
137
+
138
+ model_vocab_size = base_model.get_input_embeddings().weight.size(0)
139
+ tokenzier_vocab_size = len(tokenizer)
140
+ print(f"Vocab of the base model: {model_vocab_size}")
141
+ print(f"Vocab of the tokenizer: {tokenzier_vocab_size}")
142
+ if model_vocab_size != tokenzier_vocab_size:
143
+ assert tokenzier_vocab_size > model_vocab_size
144
+ print("Resize model embeddings to fit tokenizer")
145
+ base_model.resize_token_embeddings(tokenzier_vocab_size)
146
+ if args.lora_model is not None:
147
+ print("loading peft model")
148
+ model = PeftModel.from_pretrained(
149
+ base_model,
150
+ args.lora_model,
151
+ torch_dtype=load_type,
152
+ device_map='auto',
153
+ )
154
+ else:
155
+ model = base_model
156
+
157
+ if device == torch.device('cpu'):
158
+ model.float()
159
+
160
+ model.eval()
161
+
162
+ DEFAULT_SYSTEM_PROMPT = """
163
+ 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.
164
+
165
+ If 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.
166
+ """.strip()
167
+
168
+
169
+ def generate_prompt(prompt: str, system_prompt: str = DEFAULT_SYSTEM_PROMPT) -> str:
170
+ return f"""
171
+ [INST] <<SYS>>
172
+ {system_prompt}
173
+ <</SYS>>
174
+ {prompt} [/INST]
175
+ """.strip()
176
+
177
+ streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
178
+ text_pipeline = pipeline(
179
+ "text-generation",
180
+ model=model,
181
+ tokenizer=tokenizer,
182
+ max_new_tokens=2048,
183
+ temperature=0,
184
+ top_p=0.95,
185
+ repetition_penalty=1.15,
186
+ streamer=streamer,
187
+ )
188
+ llm = HuggingFacePipeline(pipeline=text_pipeline, model_kwargs={"temperature": 0})
189
+ SYSTEM_PROMPT = "Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer."
190
+ template = generate_prompt(
191
+ """
192
+ {context}
193
+
194
+ Question: {question}
195
+ """,
196
+ system_prompt=SYSTEM_PROMPT,
197
+ )
198
+ prompt = PromptTemplate(template=template, input_variables=["context", "question"])
199
+ qa_chain = RetrievalQA.from_chain_type(
200
+ llm=llm,
201
+ chain_type="stuff",
202
+ retriever=vector_store.as_retriever(search_kwargs={"k": 2}),
203
+ return_source_documents=True,
204
+ chain_type_kwargs={"prompt": prompt},
205
+ )
206
+
207
+ # Reset the user input
208
+ def reset_user_input():
209
+ return gr.update(value='')
210
+
211
+
212
+ # Reset the state
213
+ def reset_state():
214
+ return []
215
+
216
+
217
+ # Generate the prompt for the input of LM model
218
+ """
219
+ def generate_prompt(instruction,my_input):
220
+ return f"Instruction:{my_input}\n Response:{instruction}"
221
+ """
222
+
223
+
224
+ # User interaction function for chat
225
+ def user(user_message, history):
226
+ return gr.update(value="", interactive=False), history + \
227
+ [[user_message, None]]
228
+
229
+
230
+ class Stream(StoppingCriteria):
231
+ def __init__(self, callback_func=None):
232
+ self.callback_func = callback_func
233
+
234
+ def __call__(self, input_ids, scores) -> bool:
235
+ if self.callback_func is not None:
236
+ self.callback_func(input_ids[0])
237
+ return False
238
+
239
+
240
+ class Iteratorize:
241
+ """
242
+ Transforms a function that takes a callback
243
+ into a lazy iterator (generator).
244
+
245
+ Adapted from: https://stackoverflow.com/a/9969000
246
+ """
247
+ def __init__(self, func, kwargs=None, callback=None):
248
+ self.mfunc = func
249
+ self.c_callback = callback
250
+ self.q = Queue()
251
+ self.sentinel = object()
252
+ self.kwargs = kwargs or {}
253
+ self.stop_now = False
254
+
255
+ def _callback(val):
256
+ if self.stop_now:
257
+ raise ValueError
258
+ self.q.put(val)
259
+
260
+ def gentask():
261
+ try:
262
+ ret = self.mfunc(callback=_callback, **self.kwargs)
263
+ except ValueError:
264
+ pass
265
+ except Exception:
266
+ traceback.print_exc()
267
+
268
+ clear_torch_cache()
269
+ self.q.put(self.sentinel)
270
+ if self.c_callback:
271
+ self.c_callback(ret)
272
+
273
+ self.thread = Thread(target=gentask)
274
+ self.thread.start()
275
+
276
+ def __iter__(self):
277
+ return self
278
+
279
+ def __next__(self):
280
+ obj = self.q.get(True, None)
281
+ if obj is self.sentinel:
282
+ raise StopIteration
283
+ else:
284
+ return obj
285
+
286
+ def __del__(self):
287
+ clear_torch_cache()
288
+
289
+ def __enter__(self):
290
+ return self
291
+
292
+ def __exit__(self, exc_type, exc_val, exc_tb):
293
+ self.stop_now = True
294
+ clear_torch_cache()
295
+
296
+
297
+ def clear_torch_cache():
298
+ gc.collect()
299
+ if torch.cuda.device_count() > 0:
300
+ torch.cuda.empty_cache()
301
+
302
+
303
+ # Perform prediction based on the user input and history
304
+ @torch.no_grad()
305
+ def predict(
306
+ history,
307
+ max_new_tokens=128,
308
+ top_p=0.75,
309
+ temperature=0.1,
310
+ top_k=40,
311
+ do_sample=True,
312
+ repetition_penalty=1.0
313
+ ):
314
+ history[-1][1] = ""
315
+
316
+ history[-1][1] = qa_chain(history[-1][0])['result']
317
+ """
318
+ #history的格式:[[query1,response1],[query2,response2],[query3,response3]……]
319
+ docs=vector_store.similarity_search(history[-1][0])
320
+ context=[doc.page_content for doc in docs]
321
+ #使用下面的方式,把多轮对话转为单轮对话
322
+ input = f"### Instruction:{history[-1][0]} ### Response:{history[-1][1]}"
323
+ prompt = generate_prompt(input,"".join(context))
324
+ inputs = tokenizer(qa_chain, return_tensors="pt")
325
+ input_ids = inputs["input_ids"].to(device)
326
+
327
+ generate_params = {
328
+ 'input_ids': input_ids,
329
+ 'max_new_tokens': max_new_tokens,
330
+ 'top_p': top_p,
331
+ 'temperature': temperature,
332
+ 'top_k': top_k,
333
+ 'do_sample': do_sample,
334
+ 'repetition_penalty': repetition_penalty,
335
+ }
336
+
337
+ def generate_with_callback(callback=None, **kwargs):
338
+ if 'stopping_criteria' in kwargs:
339
+ kwargs['stopping_criteria'].append(Stream(callback_func=callback))
340
+ else:
341
+ kwargs['stopping_criteria'] = [Stream(callback_func=callback)]
342
+ clear_torch_cache()
343
+ with torch.no_grad():
344
+ model.generate(**kwargs)
345
+
346
+ def generate_with_streaming(**kwargs):
347
+ return Iteratorize(generate_with_callback, kwargs, callback=None)
348
+
349
+ with generate_with_streaming(**generate_params) as generator:
350
+ for output in generator:
351
+ next_token_ids = output[len(input_ids[0]):]
352
+ if next_token_ids[0] == tokenizer.eos_token_id:
353
+ break
354
+ new_tokens = tokenizer.decode(
355
+ next_token_ids, skip_special_tokens=True)
356
+ if isinstance(tokenizer, LlamaTokenizer) and len(next_token_ids) > 0:
357
+ if tokenizer.convert_ids_to_tokens(int(next_token_ids[0])).startswith('▁'):
358
+ new_tokens = ' ' + new_tokens
359
+
360
+ history[-1][1] = new_tokens
361
+ yield history
362
+ if len(next_token_ids) >= max_new_tokens:
363
+ break
364
+ """
365
+ yield history
366
+
367
+ # Call the setup function to initialize the components
368
+ setup()
369
+
370
+
371
+ # Create the Gradio interface
372
+ with gr.Blocks() as demo:
373
+ github_banner_path = 'https://radformation.com/images/radformation-logo-white.svg'
374
+ #gr.HTML(f'<p align="center"><a href="https://radformation.com/"><img src={github_banner_path} width="700"/></a></p>')
375
+ gr.Markdown("> Radformation Q&A bot")
376
+ chatbot = gr.Chatbot()
377
+ with gr.Row():
378
+ with gr.Column(scale=4):
379
+ with gr.Column(scale=12):
380
+ user_input = gr.Textbox(
381
+ show_label=False,
382
+ placeholder="Shift + Enter, to send message...",
383
+ lines=10).style(
384
+ container=False)
385
+ with gr.Column(min_width=32, scale=1):
386
+ submitBtn = gr.Button("Submit", variant="primary")
387
+ with gr.Column(scale=1):
388
+ emptyBtn = gr.Button("Clear History")
389
+ max_new_token = gr.Slider(
390
+ 0,
391
+ 4096,
392
+ value=512,
393
+ step=1.0,
394
+ label="Maximum New Token Length",
395
+ interactive=True)
396
+ top_p = gr.Slider(0, 1, value=0.9, step=0.01,
397
+ label="Top P", interactive=True)
398
+ temperature = gr.Slider(
399
+ 0,
400
+ 1,
401
+ value=0.5,
402
+ step=0.01,
403
+ label="Temperature",
404
+ interactive=True)
405
+ top_k = gr.Slider(1, 40, value=40, step=1,
406
+ label="Top K", interactive=True)
407
+ do_sample = gr.Checkbox(
408
+ value=True,
409
+ label="Do Sample",
410
+ info="use random sample strategy",
411
+ interactive=True)
412
+ repetition_penalty = gr.Slider(
413
+ 1.0,
414
+ 3.0,
415
+ value=1.1,
416
+ step=0.1,
417
+ label="Repetition Penalty",
418
+ interactive=True)
419
+
420
+ params = [user_input, chatbot]
421
+ predict_params = [
422
+ chatbot,
423
+ max_new_token,
424
+ top_p,
425
+ temperature,
426
+ top_k,
427
+ do_sample,
428
+ repetition_penalty]
429
+
430
+ submitBtn.click(
431
+ user,
432
+ params,
433
+ params,
434
+ queue=False).then(
435
+ predict,
436
+ predict_params,
437
+ chatbot).then(
438
+ lambda: gr.update(
439
+ interactive=True),
440
+ None,
441
+ [user_input],
442
+ queue=False)
443
+
444
+ user_input.submit(
445
+ user,
446
+ params,
447
+ params,
448
+ queue=False).then(
449
+ predict,
450
+ predict_params,
451
+ chatbot).then(
452
+ lambda: gr.update(
453
+ interactive=True),
454
+ None,
455
+ [user_input],
456
+ queue=False)
457
+
458
+ submitBtn.click(reset_user_input, [], [user_input])
459
+
460
+ emptyBtn.click(reset_state, outputs=[chatbot], show_progress=True)
461
+
462
+
463
+ # Launch the Gradio interface
464
+ """
465
+ demo.queue().launch(
466
+ share=share,
467
+ inbrowser=True,
468
+ server_name='0.0.0.0',
469
+ server_port=port)
470
+ """
471
+
472
+ demo.queue().launch(
473
+ root_path="/etc/nginx/sites-available/radllama2_gradio_app"
474
+ )