Spaces:

kaixuan42
/

aigc-lab

Runtime error

App Files Files Community

aigc-lab / app.py

kaixuan42

Update app.py

ef12332 11 months ago

raw history blame contribute delete

No virus

4.3 kB

	import json
	import torch
	import streamlit as st
	from transformers import AutoModelForCausalLM, AutoTokenizer
	from transformers.generation.utils import GenerationConfig


	st.set_page_config(page_title="Baichuan-13B-Chat")
	st.title("Baichuan-13B-Chat")


	@st.cache_resource
	def init_model():
	model = AutoModelForCausalLM.from_pretrained(
	"baichuan-inc/Baichuan-13B-Chat",
	torch_dtype=torch.float16,
	device_map="auto",
	trust_remote_code=True
	)
	model.generation_config = GenerationConfig.from_pretrained(
	"baichuan-inc/Baichuan-13B-Chat"
	)
	tokenizer = AutoTokenizer.from_pretrained(
	"baichuan-inc/Baichuan-13B-Chat",
	use_fast=False,
	trust_remote_code=True
	)
	return model, tokenizer


	# def clear_chat_history():
	# del st.session_state.messages
	# 清空聊天历史
	def clear_chat_history():
	# 把会话状态中的消息记录清空，并刷新页面
	st.session_state.messages = []
	st.experimental_rerun()

	# def init_chat_history():
	# with st.chat_message("assistant", avatar='🤖'):
	# st.markdown("您好，我是百川大模型，很高兴为您服务🥰")

	# if "messages" in st.session_state:
	# for message in st.session_state.messages:
	# avatar = '🧑‍💻' if message["role"] == "user" else '🤖'
	# with st.chat_message(message["role"], avatar=avatar):
	# st.markdown(message["content"])
	# else:
	# st.session_state.messages = []

	# return st.session_state.messages

	# 初始化聊天历史
	def init_chat_history():
	# 显示机器人的欢迎语
	st.write("🤖: 您好，我是百川大模型，很高兴为您服务🥰")

	# 如果会话状态中有消息记录，就显示出来
	if "messages" in st.session_state:
	for message in st.session_state.messages:
	# 根据角色显示不同的头像
	avatar = '🧑‍💻' if message["role"] == "user" else '🤖'
	# 用st.write或st.markdown显示消息内容
	st.write(f"{avatar}: {message['content']}")
	else:
	# 如果没有消息记录，就初始化一个空列表
	st.session_state.messages = []

	return st.session_state.messages

	# def main():
	# model, tokenizer = init_model()
	# messages = init_chat_history()

	# if prompt := st.chat_input("Shift + Enter 换行, Enter 发送"):
	# with st.chat_message("user", avatar='🧑‍💻'):
	# st.markdown(prompt)
	# messages.append({"role": "user", "content": prompt})
	# print(f"[user] {prompt}", flush=True)
	# with st.chat_message("assistant", avatar='🤖'):
	# placeholder = st.empty()
	# for response in model.chat(tokenizer, messages, stream=True):
	# placeholder.markdown(response)
	# if torch.backends.mps.is_available():
	# torch.mps.empty_cache()
	# messages.append({"role": "assistant", "content": response})
	# print(json.dumps(messages, ensure_ascii=False), flush=True)

	# st.button("清空对话", on_click=clear_chat_history)
	# 主函数
	def main():
	model, tokenizer = init_model()
	messages = init_chat_history()

	# 在侧边栏或表单中创建一个输入框和一个提交按钮
	with st.sidebar.form("chat_form"):
	user_input = st.text_input("请输入您的问题")
	submit_button = st.form_submit_button("发送")

	# 如果用户点击了发送按钮，就把用户的输入添加到消息记录中，并显示出来
	if submit_button:
	messages.append({"role": "user", "content": user_input})
	st.write(f"🧑‍💻: {user_input}")
	print(f"[user] {user_input}", flush=True)

	# 调用模型来生成回复，并添加到消息记录中，并显示出来
	reply = model.chat(tokenizer, messages, stream=False)
	messages.append({"role": "assistant", "content": reply})
	st.write(f"🤖: {reply}")
	print(json.dumps(messages, ensure_ascii=False), flush=True)

	# 如果有清空对话的按钮，就绑定清空聊天历史的函数
	if st.button("清空对话"):
	clear_chat_history()

	if __name__ == "__main__":
	main()