Update app.py

app.py

@@ -11,13 +11,7 @@ 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")
@@ -56,6 +50,25 @@ st.set_page_config(
 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'] = []
 
@@ -67,16 +80,16 @@ def get_text():
     input_text = st.text_input("用户: ","你好！", key="input")
     return input_text  
 
-user_history = []
-bot_history = []
+#user_history = []
+#bot_history = []
 user_input = get_text()
-user_history.append(user_input)
+#user_history.append(user_input)
 
 if user_input:
-    output = answer(user_history,bot_history)
     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)
+    #bot_history.append(output)
 
 if st.session_state['generated']:
 
@@ -86,4 +99,16 @@ if st.session_state['generated']:
     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'] = []