112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
|
from vector_store import get_client
|
|||
|
|
from split_markdown import split_markdown
|
|||
|
|
from embedding import get_embedding
|
|||
|
|
import gradio as gr
|
|||
|
|
import os
|
|||
|
|
import pickle
|
|||
|
|
from llm.wenxin import Wenxin, ModelName
|
|||
|
|
from dotenv import load_dotenv
|
|||
|
|
load_dotenv()
|
|||
|
|
|
|||
|
|
|
|||
|
|
chroma_client = get_client()
|
|||
|
|
collection = chroma_client.get_collection(name="amis")
|
|||
|
|
|
|||
|
|
wenxin = Wenxin()
|
|||
|
|
|
|||
|
|
text_blocks_by_id = {}
|
|||
|
|
with open(os.path.join(os.path.dirname(__file__), 'text.pickle'), 'rb') as f:
|
|||
|
|
text_blocks_by_id = pickle.load(f)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_prompt(context, query):
|
|||
|
|
return f"""
|
|||
|
|
请只根据下面的资料回答问题，如果无法根据这些资料回答，回答“找不到相关答案”：
|
|||
|
|
|
|||
|
|
资料：
|
|||
|
|
{context}
|
|||
|
|
|
|||
|
|
问题是：{query}
|
|||
|
|
回答："""
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_context(search_result, include_code=True, max_length=1024):
|
|||
|
|
context = ""
|
|||
|
|
doc_ids = []
|
|||
|
|
for doc_id in search_result['ids'][0]:
|
|||
|
|
doc_id = doc_id.split("_")[0]
|
|||
|
|
if doc_id not in doc_ids:
|
|||
|
|
doc_ids.append(doc_id)
|
|||
|
|
|
|||
|
|
for doc_id in doc_ids:
|
|||
|
|
markdown_block = text_blocks_by_id[doc_id]
|
|||
|
|
block_text = markdown_block.gen_text(512, include_code)
|
|||
|
|
if (len(context) + len(block_text)) < max_length:
|
|||
|
|
context += block_text + "\n\n"
|
|||
|
|
|
|||
|
|
return context
|
|||
|
|
|
|||
|
|
|
|||
|
|
query = gr.Textbox(label="问题")
|
|||
|
|
include_code = gr.Checkbox(value=True, label="提示词中是否要包含 amis schema",
|
|||
|
|
info="包含的好处是大模型会返回 json，但也会导致内容太长，只能提供少量段落给大模型，导致错过重要资料")
|
|||
|
|
n_result = gr.Number(
|
|||
|
|
value=10, precision=0, label="向量搜索查询返回个数")
|
|||
|
|
|
|||
|
|
|
|||
|
|
bot_result = gr.Textbox(label="文心的回答")
|
|||
|
|
bot_turbo_result = gr.Textbox(label="文心 Turbo 的回答")
|
|||
|
|
booomz_result = gr.Textbox(label="开源 BLOOMZ 的回答")
|
|||
|
|
prompt = gr.Textbox(label="提示词")
|
|||
|
|
vector_search_result = gr.Dataframe(
|
|||
|
|
label="向量相关搜索结果，这个结果只是为了辅助调试，确认是因为没找到相关内容还是大模型没能理解",
|
|||
|
|
headers=["相关段落", "所属文档"],
|
|||
|
|
datatype=["str", "str"],
|
|||
|
|
col_count=(2, "dynamic"),
|
|||
|
|
wrap=True
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def amis_search(query, n_result=10, include_code=True):
|
|||
|
|
if query.strip() == "":
|
|||
|
|
return "必须有输入", "", "", []
|
|||
|
|
|
|||
|
|
search_result = collection.query(
|
|||
|
|
query_embeddings=get_embedding(query).tolist(),
|
|||
|
|
n_results=n_result
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
context = get_context(search_result, include_code)
|
|||
|
|
|
|||
|
|
if (context == ""):
|
|||
|
|
return "检索不到相关内容", "", "", []
|
|||
|
|
|
|||
|
|
prompt = get_prompt(context, query)
|
|||
|
|
bot_result = wenxin.generate(prompt, ModelName.ERNIE_BOT)
|
|||
|
|
# bloomz_result = wenxin.generate(prompt, ModelName.BLOOMZ_7B)
|
|||
|
|
markdown_blocks = []
|
|||
|
|
index = 0
|
|||
|
|
for doc in search_result['documents'][0]:
|
|||
|
|
markdown_block = []
|
|||
|
|
markdown_block.append(doc)
|
|||
|
|
if index < len(search_result['metadatas'][0]):
|
|||
|
|
source = search_result['metadatas'][0][index]['source'].replace(
|
|||
|
|
'docs/zh-CN/', '')
|
|||
|
|
markdown_block.append(
|
|||
|
|
source)
|
|||
|
|
else:
|
|||
|
|
print("index out of range", doc)
|
|||
|
|
|
|||
|
|
markdown_blocks.append(markdown_block)
|
|||
|
|
index += 1
|
|||
|
|
|
|||
|
|
return bot_result, prompt, markdown_blocks
|
|||
|
|
|
|||
|
|
|
|||
|
|
demo = gr.Interface(amis_search, title="amis 文档问答机器人", inputs=[
|
|||
|
|
query, n_result, include_code], outputs=[bot_result, prompt, vector_search_result])
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
demo.queue(concurrency_count=10).launch(share=False, server_name="0.0.0.0")
|