import lazyllm
from lazyllm import (
OnlineEmbeddingModule, OnlineChatModule, Document, SentenceSplitter,
Retriever, Reranker, ChatPrompter, pipeline
)
# 初始化api key和提示词
api_key = 'sk-'
prompt = """
You will play the role of an AI Q&A assistant and complete a dialogue task.
In this task, you need to provide your answer based on the given context and question.
"""
# 初始化模型
embed_model = OnlineEmbeddingModule(source="siliconflow", api_key=api_key)
rerank_model = OnlineEmbeddingModule(source="siliconflow", api_key=api_key, type="rerank")
llm = OnlineChatModule(source="siliconflow", api_key=api_key)
# 定义文档管理模块,并创建节点组
doc = Document(dataset_path="/home/xxx/database", manager=False, embed=embed_model)
doc.create_node_group(name="block", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
doc.create_node_group(name="line", transform=SentenceSplitter, chunk_size=128, chunk_overlap=20, parent="block")
# 构建RAG pipeline(多路召回--重排--提示词拼接--大模型回答)
with pipeline() as ppl:
with lazyllm.parallel().sum as ppl.prl:
prl.r1 = Retriever(doc, group_name='line', similarity="cosine", topk=6, target='block')
prl.r2 = Retriever(doc, group_name='block', similarity="cosine", topk=6)
ppl.reranker = Reranker('ModuleReranker', model=rerank_model, output_format='content',
join=True) | bind(query=ppl.input)
ppl.formatter = (lambda context, query: dict(context_str=str(context), query=query)) | bind(query=ppl.input)
ppl.llm = llm.prompt(lazyllm.ChatPrompter(prompt, extra_keys=["context_str"]))
ppl.start()
query = "何为天道"
print(ppl(query))