Spaces:
Build error
Build error
| 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 --upgrade') | |
| os.system('pip install streamlit-chat') | |
| from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer | |
| import torch | |
| import streamlit as st | |
| from streamlit_chat import message | |
| # 修改colab笔记本设置为gpu，推理更快 | |
| device = torch.device('cpu') | |
| 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)") | |
| def load_model(): | |
| model = T5ForConditionalGeneration.from_pretrained("ClueAI/ChatYuan-large-v1") | |
| model.to(device) | |
| print('Model Load done!') | |
| return model | |
| 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.session_state['generated'] = [] | |
| st.session_state['past'] = [] |