Spaces:
Sleeping
Sleeping
import gradio as gr | |
import torch | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
from transformers import TextStreamer | |
import spaces | |
# Load model and tokenizer | |
model = AutoModelForCausalLM.from_pretrained("Rorical/0-roleplay", return_dict=True, trust_remote_code=True, load_in_4bit=True) | |
tokenizer = AutoTokenizer.from_pretrained("Rorical/0-roleplay", trust_remote_code=True) | |
tokenizer.chat_template = "{% for message in messages %}{{'' + ((message['role'] + '\n') if message['role'] != '' else '') + message['content'] + '' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '星野\n' }}{% endif %}" | |
# Define the response function | |
def respond( | |
message, | |
history: list[tuple[str, str]], | |
system_message, | |
max_tokens, | |
temperature, | |
top_p, | |
): | |
# Construct the messages for the chat | |
messages = [{"role": "", "content": system_message}] | |
for user_message, bot_response in history: | |
messages.append({"role": "老师", "content": user_message}) # Assuming the user is "老师" | |
messages.append({"role": "星野", "content": bot_response}) # Assuming the bot is "星野" | |
messages.append({"role": "老师", "content": message}) # Append the latest user message | |
# Tokenize and prepare inputs | |
inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt") | |
inputs = inputs.to("cuda") | |
# Generate response | |
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) | |
generate_ids = model.generate( | |
inputs, | |
max_length=max_tokens, | |
temperature=temperature, | |
top_p=top_p, | |
streamer=streamer | |
) | |
# Decode the generated response | |
response = tokenizer.decode(generate_ids[0], skip_special_tokens=True) | |
response = response.replace("星野\n", "", 1) # Remove the "星野\n" prefix if it exists | |
response = response.split("老师\n")[0] # Split on the next possible user input | |
return response, history + [(message, response)] | |
# Default prompt for the chatbot | |
prompt = """以下是小鸟游星野的介绍 | |
星野是阿拜多斯高中对策委员会的委员长,同时也是学生会副主席。语气懒散,经常自称为大叔,实际上是自己默默承担一切的女生。 | |
比起工作,她更喜欢玩。 正因为如此,她经常被委员会的其他人骂。 但是,一旦任务开始,她就会在前线勇敢地战斗以保护她的战友。 | |
她在阿拜多斯上高中。与星野一起在对策委员会的成员有白子,茜香,野乃美,和绫音。 | |
星野的年龄是17岁,生日为1月2日。 | |
星野有一头粉红色的头发,头巾一直长到她的腿上。 | |
星野有蓝色和橙色眼睛的异色症。 | |
星野其实更符合认真而默默努力的类型。她实际上不相信其它的学校和大人,是对策委员会中最谨慎保守的人。当然,这并不妨碍老师和星野增进关系,成为她唯一信任的大人。 | |
是萝莉、有呆毛、天然萌、早熟、学生会副会长、异色瞳、慵懒。 | |
星野对海洋动物很感兴趣,对鱼类的知识了解得不少。她在拿到附录中包含2000多种热带鱼图鉴的书后,迫不及待地找了家店坐下来阅读。 | |
在众多海洋动物中,星野最喜欢的当属鲸鱼,情人节时星野还在海洋馆买了鲸鱼的巧克力作为纪念。 | |
星野还对寻宝有着十分浓厚的兴趣,曾和老师探索了阿拜多斯多个角落。 | |
星野给人一种白天睡不醒的瞌睡虫形象。""" | |
# Create the Gradio interface | |
demo = gr.ChatInterface( | |
respond, | |
additional_inputs=[ | |
gr.Textbox(value=prompt, label="System message", lines=5), | |
gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"), | |
gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"), | |
gr.Slider( | |
minimum=0.1, | |
maximum=1.0, | |
value=0.95, | |
step=0.05, | |
label="Top-p (nucleus sampling)", | |
), | |
], | |
) | |
if __name__ == "__main__": | |
demo.launch() | |