File size: 2,119 Bytes
e936a3f bb97cbe |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import streamlit as st
from transformers import AutoModel, AutoTokenizer
import mdtex2html
from utils import load_model_on_gpus
st.set_page_config(page_title="ChatGLM2-6B", page_icon=":robot:")
st.header("ChatGLM2-6B")
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
# Load model on multiple GPUs
#model = load_model_on_gpus("THUDM/chatglm2-6b", num_gpus=2)
model = model.eval()
def postprocess(chat):
for i, (user, response) in enumerate(chat):
chat[i] = (mdtex2html.convert(user), mdtex2html.convert(response))
return chat
user_input = st.text_area("Input:", height=200, placeholder="Ask me anything!")
if user_input:
history = st.session_state.get('history', [])
max_length = st.slider("Max Length:", 0, 32768, 8192, 1)
top_p = st.slider("Top P:", 0.0, 1.0, 0.8, 0.01)
temperature = st.slider("Temperature:", 0.0, 1.0, 0.95, 0.01)
if 'past_key_values' not in st.session_state:
st.session_state['past_key_values'] = None
with st.spinner("Thinking..."):
response = model.generate(tokenizer.encode(user_input),
max_length=max_length,
top_p=top_p,
temperature=temperature,
return_dict_in_generate=True,
output_scores=True,
return_past_key_values=True,
past_key_values=st.session_state.past_key_values)
st.session_state.past_key_values = response.past_key_values
history.append((user_input, response.sequences[0]))
history = postprocess(history)
for user, chatbot in history:
message = f"**Human:** {user}" if user else ""
response = f"**AI:** {chatbot}" if chatbot else ""
st.markdown(message + response, unsafe_allow_html=True)
if st.button("Clear History"):
st.session_state['history'] = []
st.session_state['past_key_values'] = None |