dsRAG：用于处理非结构化数据和复杂查询的检索引擎

综合介绍

dsRAG是一款高性能检索引擎，专门用于处理非结构化数据的复杂查询。它在处理密集文本（如财务报告、法律文件和学术论文）中的挑战性查询时表现尤为出色。dsRAG采用了三种关键方法来提高性能：语义分段、上下文自动生成和相关段提取。这些方法使dsRAG在复杂的开卷问答任务中取得了显著高于传统RAG基线的准确性。此外，dsRAG支持多种配置选项，可根据用户需求进行自定义。其模块化设计允许用户轻松集成不同的组件，如向量数据库、嵌入模型和重排序器，从而实现最佳的检索效果。

与传统的RAG（Retrieval-Augmented Generation）基线相比，dsRAG在复杂的开放书本问答任务中实现了显著更高的准确率。例如，在FinanceBench基准测试中，dsRAG的准确率高达96.6%，而传统RAG基线仅为32%。dsRAG通过语义分段、自动上下文和相关段提取等关键方法，显著提升了检索性能。

功能列表

• 语义分段：利用LLM将文档分段，提高检索准确性。
• 上下文自动生成：生成包含文档和段落级上下文的块头，提高嵌入质量。
• 相关段提取：在查询时智能组合相关的文本块，生成更长的段落。
• 多种向量数据库支持：如BasicVectorDB、WeaviateVectorDB、ChromaDB等。
• 多种嵌入模型支持：如OpenAIEmbedding、CohereEmbedding等。
• 多种重排序器支持：如CohereReranker、VoyageReranker等。
• 持久化知识库：支持将知识库对象持久化到磁盘，方便后续加载和查询。
• 多种文档格式支持：支持PDF、Markdown等多种文档格式的解析和处理。

使用帮助

安装

要安装dsRAG的Python包，可以运行以下命令：

pip install dsrag

确保你有OpenAI和Cohere的API密钥，并将它们设置为环境变量。

快速开始

你可以使用create_kb_from_file函数直接从文件创建新的知识库：

from dsrag.create_kb import create_kb_from_file
file_path = "dsRAG/tests/data/levels_of_agi.pdf"
kb_id = "levels_of_agi"
kb = create_kb_from_file(kb_id, file_path)

知识库对象会自动持久化到磁盘，因此不需要显式保存。

现在，你可以通过其kb_id加载知识库，并使用query方法进行查询：

from dsrag.knowledge_base import KnowledgeBase
kb = KnowledgeBase("levels_of_agi")
search_queries = ["What are the levels of AGI?", "What is the highest level of AGI?"]
results = kb.query(search_queries)
for segment in results:
print(segment)

基本定制

你可以自定义知识库的配置，例如只使用OpenAI：

from dsrag.llm import OpenAIChatAPI
from dsrag.reranker import NoReranker
llm = OpenAIChatAPI(model='gpt-4o-mini')
reranker = NoReranker()
kb = KnowledgeBase(kb_id="levels_of_agi", reranker=reranker, auto_context_model=llm)

然后使用add_document方法添加文档：

file_path = "dsRAG/tests/data/levels_of_agi.pdf"
kb.add_document(doc_id=file_path, file_path=file_path)

架构

知识库对象接受文档（原始文本形式），并对其进行分块和嵌入，以及其他预处理操作。查询时，输入查询后，系统会进行向量数据库搜索、重排序和相关段提取，最终返回结果。

知识库对象默认是持久化的，其完整配置会以JSON文件形式保存，方便重建和更新。