chatbot_zh / app.py
scutcyr's picture
Update app.py
7dda2ce
import os
# install torch and tf
os.system('pip install transformers SentencePiece')
os.system('pip install torch')
from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer
import torch
import gradio as gr
# 下载模型
tokenizer = T5Tokenizer.from_pretrained("ClueAI/ChatYuan-large-v1")
model = T5ForConditionalGeneration.from_pretrained("ClueAI/ChatYuan-large-v1")
# 修改colab笔记本设置为gpu,推理更快
device = torch.device('cpu')
model.to(device)
print('Model Load done!')
def preprocess(text):
text = text.replace("\n", "\\n").replace("\t", "\\t")
return text
def postprocess(text):
return text.replace("\\n", "\n").replace("\\t", "\t")
def answer(text, sample=True, top_p=1, temperature=0.7):
'''sample:是否抽样。生成任务,可以设置为True;
top_p:0-1之间,生成的内容越多样
max_new_tokens=512 lost...'''
text = preprocess(text)
print('用户: '+text)
encoding = tokenizer(text=[text], truncation=True, padding=True, max_length=768, return_tensors="pt").to(device)
if not sample:
out = model.generate(**encoding, return_dict_in_generate=True, output_scores=False, max_new_tokens=512, num_beams=1, length_penalty=0.6)
else:
out = model.generate(**encoding, return_dict_in_generate=True, output_scores=False, max_new_tokens=512, do_sample=True, top_p=top_p, temperature=temperature, no_repeat_ngram_size=3)
out_text = tokenizer.batch_decode(out["sequences"], skip_special_tokens=True)
print('小元: '+postprocess(out_text[0]))
return postprocess(out_text[0])
def command_result(text):
output = answer(text)
return output
input_component = gr.Textbox(label = "输入你要对话的文本", value = "你好!")
output_component = gr.Textbox(label = "机器人回复")
examples = [["你好!"], ["请自我介绍一下!"]]
description = "这是一个使用ClueAI/ChatYuan-large-v1构建的中文文本对话聊天机器人"
gr.Interface(command_result, inputs = input_component, outputs=output_component, examples=examples, title = "👨🏻‍🎤 中文对话聊天机器人 👨🏻‍🎤", description=description).launch()