State-of-the-art RAG pipeline from D-Star AI

These details have not been verified by PyPI

Project links

Project description

dsRAG

Note: If you’re using (or planning to use) dsRAG in production, please fill out this short form telling us about your use case. This helps us prioritize new features. In return I’ll give you my personal email address, which you can use for priority email support.

What is dsRAG?

dsRAG is a retrieval engine for unstructured data. It is especially good at handling challenging queries over dense text, like financial reports, legal documents, and academic papers. dsRAG achieves substantially higher accuracy than vanilla RAG baselines on complex open-book question answering tasks. On one especially challenging benchmark, FinanceBench, dsRAG gets accurate answers 83% of the time, compared to the vanilla RAG baseline which only gets 19% of questions correct.

There are three key methods used to improve performance over vanilla RAG systems:

Semantic sectioning
AutoContext
Relevant Segment Extraction (RSE)

Semantic sectioning

Semantic sectioning uses an LLM to break a document into sections. It works by annotating the document with line numbers and then prompting an LLM to identify the starting and ending lines for each “semantically cohesive section.” These sections should be anywhere from a few paragraphs to a few pages long. The sections then get broken into smaller chunks if needed. The LLM is also prompted to generate descriptive titles for each section. These section titles get used in the contextual chunk headers created by AutoContext, which provides additional context to the ranking models (embeddings and reranker), enabling better retrieval.

AutoContext

AutoContext creates contextual chunk headers that contain document-level and section-level context, and prepends those chunk headers to the chunks prior to embedding them. This gives the embeddings a much more accurate and complete representation of the content and meaning of the text. In our testing, this feature leads to a dramatic improvement in retrieval quality. In addition to increasing the rate at which the correct information is retrieved, AutoContext also substantially reduces the rate at which irrelevant results show up in the search results. This reduces the rate at which the LLM misinterprets a piece of text in downstream chat and generation applications.

Relevant Segment Extraction

Relevant Segment Extraction (RSE) is a query-time post-processing step that takes clusters of relevant chunks and intelligently combines them into longer sections of text that we call segments. These segments provide better context to the LLM than any individual chunk can. For simple factual questions, the answer is usually contained in a single chunk; but for more complex questions, the answer usually spans a longer section of text. The goal of RSE is to intelligently identify the section(s) of text that provide the most relevant information, without being constrained to fixed length chunks.

For example, suppose you have a bunch of SEC filings in a knowledge base and you ask “What were Apple’s key financial results in the most recent fiscal year?” RSE will identify the most relevant segment as the entire “Consolidated Statement of Operations” section, which will be 5-10 chunks long. Whereas if you ask “Who is Apple’s CEO?” the most relevant segment will be identified as a single chunk that mentions “Tim Cook, CEO.”

Eval results

We've evaluated dsRAG on a couple of end-to-end RAG benchmarks.

FinanceBench

First, we have FinanceBench. This benchmark uses a corpus of a few hundred 10-Ks and 10-Qs. The queries are challenging, and often require combining multiple pieces of information. Ground truth answers are provided. Answers are graded manually on a pass/fail basis. Minor allowances for rounding errors are allowed, but other than that the answer must exactly match the ground truth answer to be considered correct.

The baseline retrieval pipeline, which uses standard chunking and top-k retrieval, achieves a score of 19% according to the paper. dsRAG, using default parameters and AutoQuery for query generation, achieves a score of 83%.

KITE

We couldn't find any other suitable end-to-end RAG benchmarks, so we decided to create our own, called KITE (Knowledge-Intensive Task Evaluation).

KITE currently consists of 4 datasets and a total of 50 questions.

AI Papers - ~100 academic papers about AI and RAG, downloaded from arXiv in PDF form.
BVP Cloud 10-Ks - 10-Ks for all companies in the Bessemer Cloud Index (~70 of them), in PDF form.
Sourcegraph Company Handbook - ~800 markdown files, with their original directory structure, downloaded from Sourcegraph's publicly accessible company handbook GitHub page.
Supreme Court Opinions - All Supreme Court opinions from Term Year 2022 (delivered from January '23 to June '23), downloaded from the official Supreme Court website in PDF form.

Ground truth answers are included with each sample. Most samples also include grading rubrics. Grading is done on a scale of 0-10 for each question, with a strong LLM doing the grading.

We tested four configurations:

Top-k retrieval (baseline)
Relevant segment extraction (RSE)
Top-k retrieval with contextual chunk headers (CCH)
CCH+RSE (dsRAG default config, minus semantic sectioning)

Testing RSE and CCH on their own, in addition to testing them together, lets us see the individual contributions of those two features.

Cohere English embeddings and the Cohere 3 English reranker were used for all configurations. LLM responses were generated with GPT-4o, and grading was also done with GPT-4o.

	Top-k	RSE	CCH+Top-k	CCH+RSE
AI Papers	4.5	7.9	4.7	7.9
BVP Cloud	2.6	4.4	6.3	7.8
Sourcegraph	5.7	6.6	5.8	9.4
Supreme Court Opinions	6.1	8.0	7.4	8.5
Average	4.72	6.73	6.04	8.42

Using CCH and RSE together leads to a dramatic improvement in performance, from 4.72 -> 8.42. Looking at the RSE and CCH+Top-k results, we can see that using each of those features individually leads to a large improvement over the baseline, with RSE appearing to be slightly more important than CCH.

To put these results in perspective, we also tested the CCH+RSE configuration with a smaller model, GPT-4o Mini. As expected, this led to a decrease in performance compared to using GPT-4o, but the difference was surprisingly small (7.95 vs. 8.42). Using CCH+RSE with GPT-4o Mini dramatically outperforms the baseline RAG pipeline even though the baseline uses a 17x more expensive LLM. This suggests that the LLM plays a much smaller role in end-to-end RAG system accuracy than the retrieval pipeline does.

	CCH+RSE (GPT-4o)	CCH+RSE (GPT-4o Mini)
AI Papers	7.9	7.0
BVP Cloud	7.8	7.9
Sourcegraph	9.4	8.5
Supreme Court Opinions	8.5	8.4
Average	8.42	7.95

Note: we did not use semantic sectioning for any of the configurations tested here. We'll evaluate that one separately once we finish some of the improvements we're working on for it. We also did not use AutoQuery, as the KITE questions are all suitable for direct use as search queries.

Tutorial

Installation

To install the python package, run

pip install dsrag

Quickstart

By default, dsRAG uses OpenAI for embeddings and AutoContext, and Cohere for reranking, so to run the code below you'll need to make sure you have API keys for those providers set as environmental variables with the following names: OPENAI_API_KEY and CO_API_KEY. If you want to run dsRAG with different models, take a look at the "Basic customization" section below.

You can create a new KnowledgeBase directly from a file using the create_kb_from_file helper function:

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)

KnowledgeBase objects persist to disk automatically, so you don't need to explicitly save it at this point.

Now you can load the KnowledgeBase by its kb_id (only necessary if you run this from a separate script) and query it using the query method:

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)

Basic customization

Now let's look at an example of how we can customize the configuration of a KnowledgeBase. In this case, we'll customize it so that it only uses OpenAI (useful if you don't have API keys for Anthropic and Cohere). To do so, we need to pass in a subclass of LLM and a subclass of Reranker. We'll use gpt-4o-mini for the LLM (this is what gets used for document and section summarization in AutoContext) and since OpenAI doesn't offer a reranker, we'll use the NoReranker class for that.

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)

Now we can add documents to this KnowledgeBase using the add_document method. Note that the add_document method takes in raw text, not files, so we'll have to extract the text from our file first. There are some utility functions for doing this in the document_parsing.py file.

from dsrag.document_parsing import extract_text_from_pdf

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

Architecture

KnowledgeBase object

A KnowledgeBase object takes in documents (in the form of raw text) and does chunking and embedding on them, along with a few other preprocessing operations. Then at query time you feed in queries and it returns the most relevant segments of text.

KnowledgeBase objects are persistent by default. The full configuration needed to reconstruct the object gets saved as a JSON file upon creation and updating.

Components

There are five key components that define the configuration of a KnowledgeBase, each of which are customizable:

VectorDB
ChunkDB
Embedding
Reranker
LLM

There are defaults for each of these components, as well as alternative options included in the repo. You can also define fully custom components by subclassing the base classes and passing in an instance of that subclass to the KnowledgeBase constructor.