84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
|
from threading import Thread
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
import gradio as gr
|
|||
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
|
|||
|
|
|
|||
|
|
model_id = "fireballoon/baichuan-vicuna-chinese-7b"
|
|||
|
|
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
|
|||
|
|
print("Running on device:", torch_device)
|
|||
|
|
print("CPU threads:", torch.get_num_threads())
|
|||
|
|
|
|||
|
|
|
|||
|
|
if torch_device == "cuda":
|
|||
|
|
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16).cuda()
|
|||
|
|
else:
|
|||
|
|
model = AutoModelForCausalLM.from_pretrained(model_id)
|
|||
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def run_generation(history, *args, **kwargs):
|
|||
|
|
# Get the model and tokenizer, and tokenize the user text.
|
|||
|
|
|
|||
|
|
instruction = "A chat between a curious user and an artificial intelligence assistant. " \
|
|||
|
|
"The assistant gives helpful, detailed, and polite answers to the user's questions."
|
|||
|
|
context = ''.join([f" USER: {turn[0].strip()} ASSISTANT: {turn[1].strip()} </s>" for turn in history[:-1]])
|
|||
|
|
prompt = instruction + context + f" USER: {history[-1][0].strip()} ASSISTANT:"
|
|||
|
|
|
|||
|
|
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
|
|||
|
|
|
|||
|
|
print()
|
|||
|
|
print(prompt)
|
|||
|
|
print('##', input_ids.size())
|
|||
|
|
|
|||
|
|
# Start generation on a separate thread, so that we don't block the UI. The text is pulled from the streamer
|
|||
|
|
# in the main thread. Adds timeout to the streamer to handle exceptions in the generation thread.
|
|||
|
|
streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
|
|||
|
|
generate_kwargs = dict(
|
|||
|
|
input_ids=input_ids,
|
|||
|
|
streamer=streamer,
|
|||
|
|
max_new_tokens=2048,
|
|||
|
|
do_sample=True,
|
|||
|
|
temperature=0.7,
|
|||
|
|
repetition_penalty=1.1,
|
|||
|
|
top_p=0.85
|
|||
|
|
)
|
|||
|
|
t = Thread(target=model.generate, kwargs=generate_kwargs)
|
|||
|
|
t.start()
|
|||
|
|
|
|||
|
|
# Pull the generated text from the streamer, and update the model output.
|
|||
|
|
history[-1][1] = ""
|
|||
|
|
|
|||
|
|
print("")
|
|||
|
|
for new_text in streamer:
|
|||
|
|
history[-1][1] += new_text
|
|||
|
|
print(new_text, end="", flush=True)
|
|||
|
|
yield history
|
|||
|
|
print('</s>')
|
|||
|
|
return history
|
|||
|
|
|
|||
|
|
|
|||
|
|
def reset_textbox():
|
|||
|
|
return gr.update(value='')
|
|||
|
|
|
|||
|
|
|
|||
|
|
with gr.Blocks() as demo:
|
|||
|
|
gr.Markdown(
|
|||
|
|
"# Baichuan Vicuna Chinese\n"
|
|||
|
|
f"[{model_id}](https://huggingface.co/{model_id})：使用中英双语sharegpt数据全参数微调的对话模型，基于baichuan-7b"
|
|||
|
|
)
|
|||
|
|
chatbot = gr.Chatbot().style(height=600)
|
|||
|
|
msg = gr.Textbox()
|
|||
|
|
clear = gr.ClearButton([msg, chatbot])
|
|||
|
|
|
|||
|
|
def user(user_message, history):
|
|||
|
|
return gr.update(value="", interactive=False), history + [[user_message, None]]
|
|||
|
|
|
|||
|
|
response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
|
|||
|
|
run_generation, chatbot, chatbot
|
|||
|
|
)
|
|||
|
|
response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)
|
|||
|
|
|
|||
|
|
demo.queue()
|
|||
|
|
demo.launch(server_name='0.0.0.0')
|