Merge pull request #851 from hwaking/patch-1

当前MAX_BOXES与MAX_TURNS不生效导致单次对话不终止显存不断增加最后显存溢出问题
pull/934/head^2
Zhengxiao Du 2023-05-05 23:23:25 +08:00 committed by GitHub
commit 467ce73f4c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 1 deletions

View File

@ -28,6 +28,8 @@ def predict(input, max_length, top_p, temperature, history=None):
with container:
if len(history) > 0:
if len(history)>MAX_BOXES:
history = history[-MAX_TURNS:]
for i, (query, response) in enumerate(history):
message(query, avatar_style="big-smile", key=str(i) + "_user")
message(response, avatar_style="bottts", key=str(i))
@ -66,4 +68,4 @@ if 'state' not in st.session_state:
if st.button("发送", key="predict"):
with st.spinner("AI正在思考请稍等........"):
# text generation
st.session_state["state"] = predict(prompt_text, max_length, top_p, temperature, st.session_state["state"])
st.session_state["state"] = predict(prompt_text, max_length, top_p, temperature, st.session_state["state"])