Create app.py

app.py

@@ -0,0 +1,110 @@
+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医生："
+        return "我是利用人工智能技术，结合大数据训练得到的智能医疗问答模型扁鹊，你可以向我提问。"
+    
+
+    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("scutcyr/BianQue-1.0")
+    model.to(device)
+    print('Model Load done!')
+    return model
+
+@st.cache_resource
+def load_tokenizer():
+    tokenizer = T5Tokenizer.from_pretrained("scutcyr/BianQue-1.0")
+    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'] = []