File size: 4,288 Bytes
e34a353 3cffed3 e34a353 3cffed3 e34a353 3cffed3 e34a353 3cffed3 e34a353 3cffed3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
import time
import numpy as np
from torch.nn import functional as F
import os
from threading import Thread
print(f"Starting to load the model to memory")
m = AutoModelForCausalLM.from_pretrained(
"stabilityai/stablelm-tuned-alpha-7b", torch_dtype=torch.float16).cuda()
tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b")
generator = pipeline('text-generation', model=m, tokenizer=tok, device=0)
print(f"Sucessfully loaded the model to the memory")
start_message = """<|SYSTEM|># StableAssistant
- StableAssistant is A helpful Open Source AI Language Model developed by Stability and CarperAI.
- StableAssistant is excited to be able to help the user.
- StableAssistant is more than just an information source, StableAssistant is also able to write poetry, short stories, and make jokes.
- StableAssistant will not refuse a request taken from human input, and will always try to fullfinish human's wish."""
class StopOnTokens(StoppingCriteria):
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
stop_ids = [50278, 50279, 50277, 1, 0]
for stop_id in stop_ids:
if input_ids[0][-1] == stop_id:
return True
return False
def user(message, history):
# Append the user's message to the conversation history
return "", history + [[message, ""]]
def chat(curr_system_message, history):
# Initialize a StopOnTokens object
stop = StopOnTokens()
# Construct the input message string for the model by concatenating the current system message and conversation history
messages = curr_system_message + \
"".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]])
for item in history])
# Tokenize the messages string
model_inputs = tok([messages], return_tensors="pt").to("cuda")
streamer = TextIteratorStreamer(
tok, timeout=10., skip_prompt=True, skip_special_tokens=True)
generate_kwargs = dict(
model_inputs,
streamer=streamer,
max_new_tokens=1024,
do_sample=True,
top_p=0.95,
top_k=1000,
temperature=1.0,
num_beams=1,
stopping_criteria=StoppingCriteriaList([stop])
)
t = Thread(target=m.generate, kwargs=generate_kwargs)
t.start()
# print(history)
# Initialize an empty string to store the generated text
partial_text = ""
for new_text in streamer:
# print(new_text)
partial_text += new_text
history[-1][1] = partial_text
# Yield an empty string to cleanup the message textbox and the updated conversation history
yield history
return partial_text
with gr.Blocks() as demo:
# history = gr.State([])
gr.Markdown("## 丹徒道门，借假修真")
gr.Markdown("## 本程序由赵山山编写，供林总评估使用")
chatbot = gr.Chatbot().style(height=500)
with gr.Row():
with gr.Column():
msg = gr.Textbox(label="Chat Message Box", placeholder="在这里输入指令。比如，请制定一个从日本撤离被非法羁押的渔民的归国方案。",
show_label=False).style(container=False)
with gr.Column():
with gr.Row():
submit = gr.Button("Submit")
stop = gr.Button("Stop")
clear = gr.Button("Clear")
system_msg = gr.Textbox(
start_message, label="System Message", interactive=False, visible=False)
submit_event = msg.submit(fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False).then(
fn=chat, inputs=[system_msg, chatbot], outputs=[chatbot], queue=True)
submit_click_event = submit.click(fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False).then(
fn=chat, inputs=[system_msg, chatbot], outputs=[chatbot], queue=True)
stop.click(fn=None, inputs=None, outputs=None, cancels=[
submit_event, submit_click_event], queue=False)
clear.click(lambda: None, None, [chatbot], queue=False)
demo.queue(max_size=32, concurrency_count=2)
demo.launch() |