From 905aa26b917755de89008b0c300a2f494738ce85 Mon Sep 17 00:00:00 2001
From: AdamBear <adamxiong@gmail.com>
Date: Fri, 17 Mar 2023 10:32:06 +0800
Subject: [PATCH] Create web_demo2.py

Add a steamlit based demo web_demo2.py for better UI.
need to install streamlit and streamlit-chat component fisrt:
pip install streamlit
pip install streamlit-chat

then run with the following:
streamlit run web_demo2.py --server.port 6006
---
 web_demo2.py | 56 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)
 create mode 100644 web_demo2.py

diff --git a/web_demo2.py b/web_demo2.py
new file mode 100644
index 0000000..6946a15
--- /dev/null
+++ b/web_demo2.py
@@ -0,0 +1,56 @@
+from transformers import AutoModel, AutoTokenizer
+import streamlit as st
+from streamlit_chat import message
+
+
+st.set_page_config(
+    page_title="ChatGLM-6b 演示",
+    page_icon=":robot:"
+)
+
+
+@st.cache_resource
+def get_model():
+    tokenizer = AutoTokenizer.from_pretrained("/THUDM/chatglm-6b", trust_remote_code=True)
+    model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
+    model = model.eval()
+    return tokenizer, model
+
+
+MAX_TURNS = 20
+MAX_BOXES = MAX_TURNS * 2
+
+
+def predict(input, history=None):
+    tokenizer, model = get_model()
+    if history is None:
+        history = []
+    response, history = model.chat(tokenizer, input, history)
+
+    #updates = []
+    for i, (query, response) in enumerate(history):
+        #updates.append("用户：" + query)
+        message(query, avatar_style="big-smile", key=str(i) + "_user")
+        #updates.append("ChatGLM-6B：" + response)
+        message(response, avatar_style="bottts", key=str(i))
+
+    # if len(updates) < MAX_BOXES:
+    #     updates = updates + [""] * (MAX_BOXES - len(updates))
+
+    return history
+
+
+# create a prompt text for the text generation
+prompt_text = st.text_area(label="用户命令输入",
+            height = 100,
+            placeholder="请在这儿输入您的命令")
+
+if 'state' not in st.session_state:
+    st.session_state['state'] = []
+
+if st.button("发送", key="predict"):
+    with st.spinner("AI正在思考，请稍等........"):
+        # text generation
+        st.session_state["state"] = predict(prompt_text, st.session_state["state"])
+
+    st.balloons()
\ No newline at end of file