Views
No views yet


1import gradio as gr
2import torch
3from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
4from threading import Thread
5from peft import PeftModel
6import time
7
8model_name_or_path = "mistralai/Mixtral-8x7B-Instruct-v0.1" # download weights from https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1
9lora_weights = "wangrongsheng/Aurora" # download weights from https://huggingface.co/wangrongsheng/Aurora
10
11tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
12model0 = AutoModelForCausalLM.from_pretrained(model_name_or_path, load_in_4bit=True, device_map="auto", torch_dtype=torch.bfloat16)
13model = PeftModel.from_pretrained(
14 model0,
15 lora_weights,
16)
17
18class StopOnTokens(StoppingCriteria):
19 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
20 stop_ids = [0,]
21 for stop_id in stop_ids:
22 if input_ids[0][-1] == stop_id:
23 return True
24 return False
25
26def convert_history_to_text(history):
27 text = ""
28 if len(history) > 1:
29 text = "<s> " + "".join(
30 [
31 "".join(
32 [
33 f"[INST]{item[0]}[/INST] {item[1]} ",
34 ]
35 )
36 for item in history[:-1]
37 ]
38 ) + "</s> "
39 text += "".join(
40 [
41 "".join(
42 [
43 f"[INST]{history[-1][0]}[/INST]",
44 ]
45 )
46 ]
47 )
48 return text
49
50def predict(message, history):
51
52 history_transformer_format = history + [[message, ""]]
53 stop = StopOnTokens()
54
55 messages = convert_history_to_text(history_transformer_format)
56
57 model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
58 streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
59 generate_kwargs = dict(
60 model_inputs,
61 streamer=streamer,
62 max_new_tokens=4096,
63 do_sample=True,
64 top_p=0.95,
65 top_k=1000,
66 temperature=1.0,
67 num_beams=1,
68 pad_token_id=tokenizer.eos_token_id,
69 stopping_criteria=StoppingCriteriaList([stop])
70 )
71 t = Thread(target=model.generate, kwargs=generate_kwargs)
72 t.start()
73
74 partial_message = ""
75 t1 = time.time()
76 count = 0
77 for new_token in streamer:
78 if new_token != '<':
79 partial_message += new_token
80 count += 1
81 yield partial_message
82 t2 = time.time()
83 speed = count/(t2-t1)
84 print("inference speed: %f tok/s" % speed)
85
86
87gr.ChatInterface(predict,chatbot=gr.Chatbot(height=600,),title="MoE").queue().launch()1@misc{wang2023auroraactivating,
2 title={Aurora:Activating Chinese chat capability for Mixtral-8x7B sparse Mixture-of-Experts through Instruction-Tuning},
3 author={Rongsheng Wang and Haoming Chen and Ruizhe Zhou and Yaofei Duan and Kunyan Cai and Han Ma and Jiaxi Cui and Jian Li and Patrick Cheong-Iao Pang and Yapeng Wang and Tao Tan},
4 year={2023},
5 eprint={2312.14557},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL}
8}