This commit is contained in:
2025-06-20 21:31:38 +08:00
parent 35c87df603
commit 0db8b55a3f
36 changed files with 4434 additions and 0 deletions

View File

@@ -0,0 +1,105 @@
"""
This code is supported by the website: https://www.guanjihuan.com
The newest version of this code is on the web page: https://www.guanjihuan.com/archives/38502
"""
import streamlit as st
st.set_page_config(
page_title="Chat",
layout='wide'
)
choose_load_method = 1 # 选择加载模型的方式
if choose_load_method == 0:
# 默认加载需要13G显存
@st.cache_resource
def load_model_chatglm3():
from transformers import AutoModel, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b-32k", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm3-6b-32k",trust_remote_code=True).half().cuda()
model = model.eval()
return model, tokenizer
model_chatglm3, tokenizer_chatglm3 = load_model_chatglm3()
elif choose_load_method == 1:
# 量化加载需要6G显存
@st.cache_resource
def load_model_chatglm3():
from transformers import AutoTokenizer, BitsAndBytesConfig, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b-32k", trust_remote_code=True)
nf4_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
)
model = AutoModelForCausalLM.from_pretrained("THUDM/chatglm3-6b-32k", trust_remote_code=True, quantization_config=nf4_config)
model = model.eval()
return model, tokenizer
model_chatglm3, tokenizer_chatglm3 = load_model_chatglm3()
elif choose_load_method == 2:
# 在CPU上加载需要25G内存对话速度会比较慢不推荐
@st.cache_resource
def load_model_chatglm3():
from transformers import AutoModel, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b-32k", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm3-6b-32k",trust_remote_code=True).float()
model = model.eval()
return model, tokenizer
model_chatglm3, tokenizer_chatglm3 = load_model_chatglm3()
with st.sidebar:
with st.expander('参数', expanded=True):
max_length = 409600
top_p = st.slider('top_p', 0.01, 1.0, step=0.01, value=0.8, key='top_p_session')
temperature = st.slider('temperature', 0.51, 1.0, step=0.01, value=0.8, key='temperature_session')
def reset_parameter():
st.session_state['top_p_session'] = 0.8
st.session_state['temperature_session'] = 0.8
reset_parameter_button = st.button('重置', on_click=reset_parameter)
prompt = st.chat_input("在这里输入您的命令")
def chat_response_chatglm3(prompt):
history, past_key_values = st.session_state.history_ChatGLM3, st.session_state.past_key_values_ChatGLM3
for response, history, past_key_values in model_chatglm3.stream_chat(tokenizer_chatglm3, prompt, history,
past_key_values=past_key_values,
max_length=max_length, top_p=top_p,
temperature=temperature,
return_past_key_values=True):
message_placeholder_chatglm3.markdown(response)
if stop_button:
break
st.session_state.ai_response.append({"role": "robot", "content": response, "avatar": "assistant"})
st.session_state.history_ChatGLM3 = history
st.session_state.past_key_values_ChatGLM3 = past_key_values
return response
def clear_all():
st.session_state.history_ChatGLM3 = []
st.session_state.past_key_values_ChatGLM3 = None
st.session_state.ai_response = []
if 'history_ChatGLM3' not in st.session_state:
st.session_state.history_ChatGLM3 = []
if 'past_key_values_ChatGLM3' not in st.session_state:
st.session_state.past_key_values_ChatGLM3 = None
if 'ai_response' not in st.session_state:
st.session_state.ai_response = []
for ai_response in st.session_state.ai_response:
with st.chat_message(ai_response["role"], avatar=ai_response.get("avatar")):
st.markdown(ai_response["content"])
prompt_placeholder = st.chat_message("user", avatar='user')
with st.chat_message("robot", avatar="assistant"):
message_placeholder_chatglm3 = st.empty()
if prompt:
prompt_placeholder.markdown(prompt)
st.session_state.ai_response.append({"role": "user", "content": prompt, "avatar": 'user'})
stop = st.empty()
stop_button = stop.button('停止', key='break_response')
chat_response_chatglm3(prompt)
stop.empty()
button_clear = st.button("清空", on_click=clear_all, key='clear')

View File

@@ -0,0 +1,33 @@
# basic requirements
protobuf>=4.25.2
transformers>=4.36.2
tokenizers>=0.15.0
cpm_kernels>=1.0.11
torch>=2.1.0
gradio>=4.14.0
sentencepiece>=0.1.99
sentence_transformers>=2.2.2
accelerate>=0.26.1
streamlit>=1.30.0
fastapi>=0.109.0
loguru~=0.7.2
mdtex2html>=1.2.0
latex2mathml>=3.77.0
# for openai demo
openai>=1.7.2
zhipuai>=2.0.0
pydantic>=2.5.3
sse-starlette>=1.8.2
uvicorn>=0.25.0
timm>=0.9.12
tiktoken>=0.5.2
# for langchain demo
langchain>=0.1.0
langchainhub>=0.1.14
arxiv>=2.1.0