93 lines
3.9 KiB
Python
93 lines
3.9 KiB
Python
|
import os
|
|||
|
|
from openai import OpenAI
|
|||
|
|
from llama_index.core import StorageContext,load_index_from_storage,Settings
|
|||
|
|
from llama_index.embeddings.dashscope import (
|
|||
|
|
DashScopeEmbedding,
|
|||
|
|
DashScopeTextEmbeddingModels,
|
|||
|
|
DashScopeTextEmbeddingType,
|
|||
|
|
)
|
|||
|
|
from llama_index.postprocessor.dashscope_rerank import DashScopeRerank
|
|||
|
|
from create_kb import *
|
|||
|
|
DB_PATH = "VectorStore"
|
|||
|
|
TMP_NAME = "tmp_abcd"
|
|||
|
|
EMBED_MODEL = DashScopeEmbedding(
|
|||
|
|
model_name=DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V2,
|
|||
|
|
text_type=DashScopeTextEmbeddingType.TEXT_TYPE_DOCUMENT,
|
|||
|
|
)
|
|||
|
|
# 若使用本地嵌入模型，请取消以下注释：
|
|||
|
|
# from langchain_community.embeddings import ModelScopeEmbeddings
|
|||
|
|
# from llama_index.embeddings.langchain import LangchainEmbedding
|
|||
|
|
# embeddings = ModelScopeEmbeddings(model_id="modelscope/iic/nlp_gte_sentence-embedding_chinese-large")
|
|||
|
|
# EMBED_MODEL = LangchainEmbedding(embeddings)
|
|||
|
|
|
|||
|
|
# 设置嵌入模型
|
|||
|
|
Settings.embed_model = EMBED_MODEL
|
|||
|
|
|
|||
|
|
def get_model_response(multi_modal_input,history,model,temperature,max_tokens,history_round,db_name,similarity_threshold,chunk_cnt):
|
|||
|
|
# prompt = multi_modal_input['text']
|
|||
|
|
prompt = history[-1][0]
|
|||
|
|
tmp_files = multi_modal_input['files']
|
|||
|
|
if os.path.exists(os.path.join("File",TMP_NAME)):
|
|||
|
|
db_name = TMP_NAME
|
|||
|
|
else:
|
|||
|
|
if tmp_files:
|
|||
|
|
create_tmp_kb(tmp_files)
|
|||
|
|
db_name = TMP_NAME
|
|||
|
|
# 获取index
|
|||
|
|
print(f"prompt:{prompt},tmp_files:{tmp_files},db_name:{db_name}")
|
|||
|
|
try:
|
|||
|
|
dashscope_rerank = DashScopeRerank(top_n=chunk_cnt,return_documents=True)
|
|||
|
|
storage_context = StorageContext.from_defaults(
|
|||
|
|
persist_dir=os.path.join(DB_PATH,db_name)
|
|||
|
|
)
|
|||
|
|
index = load_index_from_storage(storage_context)
|
|||
|
|
print("index获取完成")
|
|||
|
|
retriever_engine = index.as_retriever(
|
|||
|
|
similarity_top_k=20,
|
|||
|
|
)
|
|||
|
|
# 获取chunk
|
|||
|
|
retrieve_chunk = retriever_engine.retrieve(prompt)
|
|||
|
|
print(f"原始chunk为：{retrieve_chunk}")
|
|||
|
|
try:
|
|||
|
|
results = dashscope_rerank.postprocess_nodes(retrieve_chunk, query_str=prompt)
|
|||
|
|
print(f"rerank成功，重排后的chunk为：{results}")
|
|||
|
|
except:
|
|||
|
|
results = retrieve_chunk[:chunk_cnt]
|
|||
|
|
print(f"rerank失败，chunk为：{results}")
|
|||
|
|
chunk_text = ""
|
|||
|
|
chunk_show = ""
|
|||
|
|
for i in range(len(results)):
|
|||
|
|
if results[i].score >= similarity_threshold:
|
|||
|
|
chunk_text = chunk_text + f"## {i+1}:\n {results[i].text}\n"
|
|||
|
|
chunk_show = chunk_show + f"## {i+1}:\n {results[i].text}\nscore: {round(results[i].score,2)}\n"
|
|||
|
|
print(f"已获取chunk：{chunk_text}")
|
|||
|
|
prompt_template = f"请参考以下内容：{chunk_text}，以合适的语气回答用户的问题：{prompt}。如果参考内容中有图片链接也请直接返回。"
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"异常信息：{e}")
|
|||
|
|
prompt_template = prompt
|
|||
|
|
chunk_show = ""
|
|||
|
|
history[-1][-1] = ""
|
|||
|
|
client = OpenAI(
|
|||
|
|
api_key=os.getenv("DASHSCOPE_API_KEY"),
|
|||
|
|
base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
|
|||
|
|
)
|
|||
|
|
system_message = {'role': 'system', 'content': 'You are a helpful assistant.'}
|
|||
|
|
messages = []
|
|||
|
|
history_round = min(len(history),history_round)
|
|||
|
|
for i in range(history_round):
|
|||
|
|
messages.append({'role': 'user', 'content': history[-history_round+i][0]})
|
|||
|
|
messages.append({'role': 'assistant', 'content': history[-history_round+i][1]})
|
|||
|
|
messages.append({'role': 'user', 'content': prompt_template})
|
|||
|
|
messages = [system_message] + messages
|
|||
|
|
completion = client.chat.completions.create(
|
|||
|
|
model=model,
|
|||
|
|
messages=messages,
|
|||
|
|
temperature=temperature,
|
|||
|
|
max_tokens=max_tokens,
|
|||
|
|
stream=True
|
|||
|
|
)
|
|||
|
|
assistant_response = ""
|
|||
|
|
for chunk in completion:
|
|||
|
|
assistant_response += chunk.choices[0].delta.content
|
|||
|
|
history[-1][-1] = assistant_response
|
|||
|
|
yield history,chunk_show
|