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 # 下载模型 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(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)") 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: output = answer(user_history,bot_history) st.session_state.past.append(user_input) 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))