Update streamlit web demo
This commit is contained in:
parent
385040e6c7
commit
1c144d168b
13
README.md
13
README.md
@ -180,7 +180,18 @@ python web_demo.py
|
||||
```
|
||||
|
||||
程序会运行一个 Web Server,并输出地址。在浏览器中打开输出的地址即可使用。
|
||||
> 由于国内 Gradio 的网络访问较为缓慢,启用 `demo.queue().launch(share=True, inbrowser=True)` 时所有网络会经过 Gradio 服务器转发,导致打字机体验大幅下降,现在默认启动方式已经改为 `share=False`,如有需要公网访问的需求,可以重新修改为 `share=True` 启动。
|
||||
> 默认使用了 `share=False` 启动,不会生成公网链接。如有需要公网访问的需求,可以修改为 `share=True` 启动。
|
||||
>
|
||||
|
||||
感谢 [@AdamBear](https://github.com/AdamBear) 实现了基于 Streamlit 的网页版 Demo `web_demo2.py`。使用时首先需要额外安装以下依赖:
|
||||
```shell
|
||||
pip install streamlit streamlit-chat
|
||||
```
|
||||
然后通过以下命令运行:
|
||||
```shell
|
||||
streamlit run web_demo2.py
|
||||
```
|
||||
经测试,如果输入的 prompt 较长的话,使用基于 Streamlit 的网页版 Demo 会更流畅。
|
||||
|
||||
### 命令行 Demo
|
||||
|
||||
|
72
web_demo2.py
Normal file
72
web_demo2.py
Normal file
@ -0,0 +1,72 @@
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
import streamlit as st
|
||||
from streamlit_chat import message
|
||||
|
||||
|
||||
st.set_page_config(
|
||||
page_title="ChatGLM2-6b 演示",
|
||||
page_icon=":robot:",
|
||||
layout='wide'
|
||||
)
|
||||
|
||||
|
||||
@st.cache_resource
|
||||
def get_model():
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
|
||||
model = model.eval()
|
||||
return tokenizer, model
|
||||
|
||||
|
||||
MAX_TURNS = 20
|
||||
MAX_BOXES = MAX_TURNS * 2
|
||||
|
||||
|
||||
def predict(input, max_length, top_p, temperature, history=None):
|
||||
tokenizer, model = get_model()
|
||||
if history is None:
|
||||
history = []
|
||||
|
||||
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))
|
||||
|
||||
message(input, avatar_style="big-smile", key=str(len(history)) + "_user")
|
||||
st.write("AI正在回复:")
|
||||
with st.empty():
|
||||
for response, history in model.stream_chat(tokenizer, input, history, max_length=max_length, top_p=top_p,
|
||||
temperature=temperature):
|
||||
query, response = history[-1]
|
||||
st.write(response)
|
||||
|
||||
return history
|
||||
|
||||
|
||||
container = st.container()
|
||||
|
||||
# create a prompt text for the text generation
|
||||
prompt_text = st.text_area(label="用户命令输入",
|
||||
height = 100,
|
||||
placeholder="请在这儿输入您的命令")
|
||||
|
||||
max_length = st.sidebar.slider(
|
||||
'max_length', 0, 32768, 8192, step=1
|
||||
)
|
||||
top_p = st.sidebar.slider(
|
||||
'top_p', 0.0, 1.0, 0.8, step=0.01
|
||||
)
|
||||
temperature = st.sidebar.slider(
|
||||
'temperature', 0.0, 1.0, 0.95, step=0.01
|
||||
)
|
||||
|
||||
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, max_length, top_p, temperature, st.session_state["state"])
|
Loading…
Reference in New Issue
Block a user