-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathapp.py
208 lines (185 loc) · 9.03 KB
/
app.py
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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
import hydra
from hydra.core.global_hydra import GlobalHydra
from omegaconf import DictConfig, OmegaConf
import streamlit as st
from PIL import Image
import os
import sys
sys.path.append(os.path.dirname(__file__))
from download_models import download_model
@st.cache_resource
def load_simple_rag(config, used_lmdeploy=False):
## load config
data_source_dir = config["data_source_dir"]
db_persist_directory = config["db_persist_directory"]
llm_model = config["llm_model"]
embeddings_model = config["embeddings_model"]
reranker_model = config["reranker_model"]
llm_system_prompt = config["llm_system_prompt"]
rag_prompt_template = config["rag_prompt_template"]
from rag.simple_rag import WuleRAG
if not used_lmdeploy:
from rag.simple_rag import InternLM, WuleRAG
base_mode = InternLM(model_path=llm_model, llm_system_prompt=llm_system_prompt)
else:
from deploy.lmdeploy_model import LmdeployLM, GenerationConfig
cache_max_entry_count = config.get("cache_max_entry_count", 0.2)
base_mode = LmdeployLM(model_path=llm_model, llm_system_prompt=llm_system_prompt, cache_max_entry_count=cache_max_entry_count)
## loda final rag model
wulewule_rag = WuleRAG(data_source_dir, db_persist_directory, base_mode, embeddings_model, reranker_model, rag_prompt_template)
return wulewule_rag
@st.cache_resource
def load_wulewule_agent(config):
from agent.wulewule_agent import MultiModalAssistant, Settings
use_remote = config["use_remote"]
SiliconFlow_api = config["SiliconFlow_api"]
data_source_dir = config["data_source_dir"]
if len(SiliconFlow_api)<51 and os.environ.get('SiliconFlow_api', ""):
SiliconFlow_api = os.environ.get('SiliconFlow_api')
print(f"======= loading llm =======")
if use_remote:
from llama_index.llms.siliconflow import SiliconFlow
from llama_index.embeddings.siliconflow import SiliconFlowEmbedding
api_base_url = "https://api.siliconflow.cn/v1/chat/completions"
# model = "Qwen/Qwen2.5-72B-Instruct"
# model = "deepseek-ai/DeepSeek-V2.5"
remote_llm = config["remote_llm"]
remote_embeddings_model = config["remote_embeddings_model"]
llm = SiliconFlow( model=remote_llm, base_url=api_base_url, api_key=SiliconFlow_api, max_tokens=4096)
embed_model = SiliconFlowEmbedding( model=remote_embeddings_model, api_key=SiliconFlow_api)
else:
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.huggingface import HuggingFaceLLM
local_llm = config["llm_model"]
local_embeddings_model = config["agent_embeddings_model"]
llm = HuggingFaceLLM(
model_name=local_llm,
tokenizer_name=local_llm,
model_kwargs={"trust_remote_code":True},
tokenizer_kwargs={"trust_remote_code":True},
# context_window=4096,
# max_new_tokens=4096,
)
embed_model = HuggingFaceEmbedding(
model_name=local_embeddings_model
)
# settings
Settings.llm = llm
Settings.embed_model = embed_model
wulewule_assistant = MultiModalAssistant(data_source_dir, llm, SiliconFlow_api)
print(f"======= finished loading ! =======")
return wulewule_assistant
GlobalHydra.instance().clear()
@hydra.main(version_base=None, config_path="./configs", config_name="model_cfg")
def main(cfg):
# omegaconf.dictcfg.DictConfig 转换为普通字典
config_dict = OmegaConf.to_container(cfg, resolve=True)
## download model from modelscope
if not config_dict["use_remote"] and not os.path.exists(config_dict["llm_model"]):
download_model(llm_model_path =config_dict["llm_model"])
## agent mode, used llama-index, rturn off lmdeloy and chroma rag
if cfg.agent_mode:
## load wulewule agent
wulewule_assistant = load_wulewule_agent(config_dict)
cfg.use_rag = False
cfg.use_lmdepoly = False
if cfg.use_rag:
## load rag model
wulewule_model = load_simple_rag(config_dict, used_lmdeploy=cfg.use_lmdepoly)
elif ( cfg.use_lmdepoly):
## load lmdeploy model
from deploy.lmdeploy_model import load_turbomind_model, GenerationConfig
wulewule_model = load_turbomind_model(config_dict["llm_model"], config_dict["llm_system_prompt"], config_dict["cache_max_entry_count"])
## streamlit setting
if "messages" not in st.session_state:
st.session_state["messages"] = []
# 在侧边栏中创建一个标题和一个链接
with st.sidebar:
st.markdown("## 悟了悟了💡")
logo_path = "assets/sd_wulewule.webp"
if os.path.exists(logo_path):
image = Image.open(logo_path)
st.image(image, caption='wulewule')
"[InternLM](https://github.com/InternLM)"
"[悟了悟了](https://github.com/xzyun2011/wulewule.git)"
# 创建一个标题
st.title("悟了悟了：黑神话悟空AI助手🐒")
# 遍历session_state中的所有消息，并显示在聊天界面上
for msg in st.session_state.messages:
st.chat_message("user").write(msg["user"])
assistant_res = msg["assistant"]
if isinstance(assistant_res, str):
st.chat_message("assistant").write(assistant_res)
elif cfg.agent_mode and isinstance(assistant_res, dict):
image_url = assistant_res["image_url"]
audio_text = assistant_res["audio_text"]
st.chat_message("assistant").write(assistant_res["response"])
if image_url:
# 使用st.image展示URL图像，并设置使用列宽
st.image( image_url, width=256 )
if audio_text:
# 使用st.audio函数播放音频
st.audio("audio.mp3")
st.write(f"语音内容为: \n\n{audio_text}")
# Get user input
if prompt := st.chat_input("请输入你的问题，换行使用Shfit+Enter。"):
# Display user input
st.chat_message("user").write(prompt)
## 初始化完整的回答字符串
full_answer = ""
if cfg.agent_mode:
with st.chat_message('robot'):
message_placeholder = st.empty()
response_dict = wulewule_assistant.chat(prompt)
image_url = response_dict["image_url"]
audio_text = response_dict["audio_text"]
for cur_response in response_dict["response"]:
full_answer += cur_response
# Display robot response in chat message container
message_placeholder.markdown(full_answer + '▌')
message_placeholder.markdown(full_answer)
# 将问答结果添加到 session_state 的消息历史中
st.session_state.messages.append({"user": prompt, "assistant": response_dict})
if image_url:
# 使用st.image展示URL图像，并设置使用列宽
st.image( image_url, width=256 )
if audio_text:
# 使用st.audio函数播放音频
st.audio("audio.mp3")
st.write(f"语音内容为: \n\n{audio_text}")
# 流式显示, used streaming result
else:
if cfg.stream_response:
# rag
with st.chat_message('robot'):
message_placeholder = st.empty()
if cfg.use_rag:
for cur_response in wulewule_model.query_stream(prompt):
full_answer += cur_response
# Display robot response in chat message container
message_placeholder.markdown(full_answer + '▌')
elif cfg.use_lmdepoly:
# gen_config = GenerationConfig(top_p=0.8,
# top_k=40,
# temperature=0.8,
# max_new_tokens=2048,
# repetition_penalty=1.05)
messages = [{'role': 'user', 'content': f'{prompt}'}]
for response in wulewule_model.stream_infer(messages):
full_answer += response.text
# Display robot response in chat message container
message_placeholder.markdown(full_answer + '▌')
message_placeholder.markdown(full_answer)
# 一次性显示结果
else:
if cfg.use_lmdepoly:
messages = [{'role': 'user', 'content': f'{prompt}'}]
full_answer = wulewule_model(messages).text
elif cfg.use_rag:
full_answer = wulewule_model.query(prompt)
# 显示回答
st.chat_message("assistant").write(full_answer)
# 将问答结果添加到 session_state 的消息历史中
st.session_state.messages.append({"user": prompt, "assistant": full_answer})
if __name__ == "__main__":
main()