Spaces:
Build error
Build error
File size: 3,909 Bytes
7646d66 4a4d251 7646d66 e917743 7646d66 df8d5c1 7646d66 e917743 7646d66 4a4d251 7646d66 4a4d251 7646d66 4a4d251 7646d66 4a4d251 7646d66 4a4d251 7646d66 30d1742 7646d66 4a4d251 |
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 104 105 106 107 108 109 110 111 112 113 114 |
import os
# install torch and tf
os.system('pip install transformers SentencePiece')
os.system('pip install torch')
# pip install streamlit-chat
os.system('pip install streamlit-chat')
from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer
import torch
import streamlit as st
from streamlit_chat import message
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(user_history, bot_history, sample=True, top_p=1, temperature=0.7):
'''sample：是否抽样。生成任务，可以设置为True;
top_p：0-1之间，生成的内容越多样
max_new_tokens=512 lost...'''
if len(bot_history)>0:
context = "\n".join([f"用户：{user_history[i]}\n小元：{bot_history[i]}" for i in range(len(bot_history))])
input_text = context + "\n用户：" + user_history[-1] + "\n小元："
else:
input_text = "用户：" + user_history[-1] + "\n小元："
input_text = preprocess(input_text)
print(input_text)
encoding = tokenizer(text=input_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])
st.set_page_config(
page_title="Chinese ChatBot - Demo",
page_icon=":robot:"
)
st.header("Chinese ChatBot - Demo")
st.markdown("[Github](https://github.com/scutcyr)")
@st.cache_resource
def load_model():
model = T5ForConditionalGeneration.from_pretrained("ClueAI/ChatYuan-large-v1")
# 修改colab笔记本设置为gpu，推理更快
device = torch.device('cpu')
model.to(device)
print('Model Load done!')
return model
@st.cache_resource
def load_tokenizer():
tokenizer = T5Tokenizer.from_pretrained("ClueAI/ChatYuan-large-v1")
print('Tokenizer Load done!')
return tokenizer
model = load_model()
tokenizer = load_tokenizer()
if 'generated' not in st.session_state:
st.session_state['generated'] = []
if 'past' not in st.session_state:
st.session_state['past'] = []
def get_text():
input_text = st.text_input("用户: ","你好！", key="input")
return input_text
#user_history = []
#bot_history = []
user_input = get_text()
#user_history.append(user_input)
if user_input:
st.session_state.past.append(user_input)
output = answer(st.session_state['past'],st.session_state["generated"])
st.session_state.generated.append(output)
#bot_history.append(output)
if st.session_state['generated']:
#for i in range(len(st.session_state['generated'])-1, -1, -1):
# message(st.session_state["generated"][i], key=str(i))
# message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')
for i in range(len(st.session_state['generated'])):
message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')
message(st.session_state["generated"][i], key=str(i))
if st.button("清理模型缓存"):
# Clear values from *all* all in-memory and on-disk data caches:
# i.e. clear values from both square and cube
st.cache_resource.clear()
torch.cuda.empty_cache()
if st.button("清理对话缓存"):
# Clear values from *all* all in-memory and on-disk data caches:
# i.e. clear values from both square and cube
st.session_state['generated'] = []
st.session_state['past'] = [] |