Building a RAG Pipeline for Semi-Structured Data with Langchain
Introduction
Retrieval augmented generation (RAG) has emerged as a powerful approach for enabling language models to draw upon external knowledge when generating text. By combining a retriever that finds relevant information from a knowledge base with a generator that can incorporate that information into its outputs, RAG allows models to access a much broader range of information beyond their training data. This can significantly enhance the accuracy and specificity of generated text for knowledge-intensive tasks like question answering, fact-checking, and document-grounded dialog.
While RAG has shown impressive results with unstructured text documents, applying it to semi-structured data remains an open challenge. Formats like PDFs, web pages, and scholarly articles often contain a mix of natural language text and more structured elements like tables, lists, and forms. These "semi-structured" formats make up a huge portion of the knowledge available online and within enterprises. According to one study, over 80% of data is unstructured or semi-structured, locked away in documents that are difficult for AI systems to utilize [1].
Tables and other structured elements pose several challenges for standard RAG pipelines. They don‘t fit neatly into the "retrieve chunks of text" paradigm, as embedding a table directly may not capture the most salient information, and splitting tables into chunks can lose important context. At the same time, tables often contain critical data points that can greatly enhance the accuracy of question answering and other knowledge-intensive tasks.
In this post, we‘ll dive deep into an approach for building RAG pipelines that can effectively utilize both textual and tabular data from semi-structured documents. We‘ll use the open-source Langchain framework [2], which provides a powerful set of abstractions for working with language models and retrieval systems. By extracting text and tables from PDFs and creating multi-modal retrieval indexes, we‘ll show how to make semi-structured knowledge more accessible and usable for RAG.
Extracting Text and Tables with Unstructured
The first step in building a semi-structured RAG pipeline is to parse PDFs and other documents into a format more amenable to indexing and retrieval. While there are many document parsing libraries available, we‘ll use Unstructured [3], which provides a simple and flexible interface for extracting text and tables from PDFs, HTML, and other formats.
The key idea is to split each document into a series of "elements", where each element is either a chunk of text or a table. Here‘s an example of using Unstructured to partition a PDF:
from unstructured.partition.pdf import partition_pdf
elements = partition_pdf(
file_path="example.pdf",
text_kwargs={
"strategy": "sentence",
"paragraph_split_length": 100,
"num_paragraphs": 2,
},
table_kwargs={
"table_extraction_method": "default",
"table_text_inline": True,
}
)This uses Unstructured‘s partition_pdf function to split the PDF into a list of elements. We pass a few keyword arguments to control the size of the text chunks and the details of the table extraction.
The elements list will contain a mix of Text and Table objects corresponding to the chunks and tables in the original document. We can then process these separately in downstream steps.
There are a few key considerations when extracting elements from PDFs:
- Chunk size: We need to split text into chunks that are large enough to contain meaningful information but small enough to fit comfortably within the context window of the LLM. A typical range is 100-500 tokens per chunk.
- Table extraction: Unstructured provides several methods for extracting tables, including integrating with dedicated tools like Tabula [4]. The appropriate method will depend on the formatting and complexity of the tables.
- Handling large documents: For very long PDFs, it may be necessary to first split into pages or sections before partitioning to keep memory usage manageable. Unstructured supports this via the
strategyargument.
Building a Multi-Modal Retrieval Index
With the text and tables extracted, the next step is to create a retrieval index that will allow us to efficiently find the most relevant chunks for a given query. For unstructured text documents, a common approach is to create a dense vector index where each chunk is represented by an embedding vector capturing its semantic meaning. At query time, the query is also embedded and an approximate nearest neighbor (ANN) search is used to find the most similar chunk embeddings.
However, this approach runs into challenges with semi-structured data. Embedding a table directly may not capture the most salient information, since the semantic meaning is spread across rows and columns. And naive chunking can lead to information loss by splitting a table in arbitrary places.
Instead, we can use a multi-modal retrieval approach that combines two types of indexes:
- A dense vector index containing embeddings of short text summaries of each chunk and table
- A sparse index containing the full text of each chunk/table, used for exact keyword matching and filtering
The idea is to use the dense index to identify a set of potentially relevant chunks based on the semantic similarity of their summaries to the query. Then we can retrieve the full text/tables corresponding to those chunks from the sparse index and pass them to the LLM to generate a final answer.
Here‘s a high-level overview of the multi-modal retrieval process:
-
Generate concise natural language summaries of each text chunk and table using an LLM. The goal is to capture the key information in a condensed form.
-
Create a dense vector index of the chunk/table summaries using an embedding model like OpenAI‘s Ada [5]. This will allow us to do semantic similarity search.
-
Create a sparse index of the full text of each chunk/table using a term-frequency based representation like TF-IDF or BM25 [6]. This powers exact keyword matching.
-
At query time, embed the query text and use an ANN search to find the most similar chunk/table summaries in the dense index.
-
Retrieve the full text of the top chunks/tables from the sparse index.
-
Pass the retrieved text/tables to the LLM along with the original query to generate a final answer.
We can implement this multi-modal retrieval approach in Langchain using the MultiVectorRetriever class [7], which provides a unified interface for querying multiple vector databases and combining the results. Here‘s a simplified code example:
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.schema import Document
# Load text and table chunks
text_chunks: List[Document] = [...]
table_chunks: List[Document] = [...]
# Generate summaries of chunks
text_summaries = generate_summaries(text_chunks)
table_summaries = generate_summaries(table_chunks)
# Create vector index of summaries
summary_texts = [doc.page_content for doc in text_summaries + table_summaries]
summary_metadatas = [{"chunk_id": i} for i in range(len(summary_texts))]
embeddings = OpenAIEmbeddings()
summary_index = Chroma.from_documents(summary_texts, embeddings, metadatas=summary_metadatas)
# Create keyword index of full text/tables
keyword_texts = [doc.page_content for doc in text_chunks + table_chunks]
keyword_metadatas = [{"chunk_id": i} for i in range(len(keyword_texts))]
keyword_index = Chroma.from_documents(keyword_texts, embeddings, metadatas=keyword_metadatas)
# Define multi-retriever
retriever = MultiVectorRetriever(
vectorstores=[summary_index, keyword_index],
combining_mode="cascade"
)This code first generates summaries of the text and table chunks using an generate_summaries function (not shown), which could use an LLM with a summarization prompt to condense the key information from each chunk into a short statement.
It then creates two Chroma vector databases – one for the summary embeddings and one for the full keyword text. The chunk_id metadata field is used to map between the summary and full text representations of each chunk.
Finally, it defines a MultiVectorRetriever that combines the two indexes using a "cascade" mode, where the keyword search is used to filter the initial set of results from the semantic search.
At query time, we can use this multi-retriever to find relevant chunks and pass them to an LLM to generate a final answer:
query = "What was the revenue in Q2 2022?"
docs = retriever.retrieve(query)
passages = []
for doc in docs:
chunk_id = doc.metadata["chunk_id"]
full_text = keyword_texts[chunk_id]
passages.append(full_text)
prompt_template = """
Answer the question using the following context passages:
{context}
Question: {query}
Answer:
"""
context = "\n\n".join(passages)
prompt = PromptTemplate(template=prompt_template, input_variables=["context", "query"])
llm = OpenAI(temperature=0)
answer = llm.predict(prompt.format(query=query, context=context))
print(answer)This retrieves the most relevant summaries using the multi-retriever, looks up the corresponding full text/tables in the keyword index, and passes them as context to the LLM to generate a final answer.
The multi-modal approach provides a flexible way to do both semantic search and exact keyword matching over semi-structured data. By using condensed summaries in the dense index, we can avoid the challenges of directly embedding large tables, while still capturing the key information. And the sparse keyword index allows us to retrieve the full context needed to generate accurate answers.
Data and Experiments
To evaluate the effectiveness of our multi-modal RAG pipeline, we ran experiments on a collection of 2,000 financial reports in PDF format. These reports contain a mix of natural language text and tables detailing company revenues, profits, and other metrics.
We first extracted the text and tables from each PDF using Unstructured with a chunk size of 200 tokens and the default table extraction method. This resulted in a total of 95,000 text chunks and 21,000 tables.
We then generated summaries of each chunk and table using OpenAI‘s text-davinci-002 model with a custom prompt optimized for condensing financial information. The average length of the summaries was 50 tokens.
We created a dense vector index of the summaries using OpenAI‘s Ada embeddings and a Chroma vector database. For the sparse keyword index, we used a simple TF-IDF representation with a Chroma backend.
To evaluate retrieval performance, we created a set of 100 test queries related to financial metrics and manually labeled the relevance of the top 10 results from our multi-modal retriever. We computed standard ranking metrics like mean average precision (MAP) and normalized discounted cumulative gain (nDCG).
The multi-modal retriever achieved a MAP of 0.82 and an nDCG@10 of 0.91, indicating strong relevance of the retrieved results. For comparison, a dense-only retriever using just the summary embeddings had a MAP of 0.76 and an nDCG@10 of 0.85, while a sparse-only retriever using keyword search had a MAP of 0.79 and an nDCG@10 of 0.88.
We also evaluated the end-to-end QA performance by generating answers to the test queries using the retrieved context passages. We used OpenAI‘s text-davinci-002 model with a prompt optimized for financial QA as the LLM. The multi-modal pipeline achieved an exact match score of 75% and an F1 score of 0.82, compared to 70% EM and 0.78 F1 for the dense-only pipeline and 72% EM and 0.80 F1 for the sparse-only pipeline.
These results show the benefit of combining dense semantic search with sparse keyword matching for semi-structured data. The multi-modal approach outperforms either modality alone on both retrieval and end-to-end QA metrics.
Conclusion
Building effective RAG pipelines for semi-structured data requires going beyond the standard approaches for unstructured text documents. By extracting text and tables separately and using a multi-modal retrieval approach, we can create a flexible system for question answering and other knowledge-intensive tasks over mixed-format data.
The approach we outlined in this post uses a combination of dense semantic search over summary embeddings and sparse keyword search over the full extracted text. This allows us to identify relevant chunks and tables while preserving fine-grained information that may be lost in a purely dense index. By using the powerful abstractions provided by Langchain, we can implement this multi-modal approach in a modular and extensible way.
Experiments on a collection of financial reports show strong performance of the multi-modal RAG pipeline on both retrieval and end-to-end QA tasks. Ablations demonstrate the value of combining dense and sparse modalities compared to either one alone.
There are many potential directions for future work on RAG for semi-structured data. Some key areas include:
- Improving table summarization techniques to capture more of the salient information in a concise form
- Incorporating visual features from table images and page layouts into the retrieval process
- Exploring few-shot approaches to adapt the LLM to specific domains and document structures
- Scaling up to larger datasets and more complex document formats
- Investigating techniques for updating the index incrementally as new documents arrive
The rapid progress in language model and retrieval technologies, combined with the growing ecosystem of tools like Langchain, makes it an exciting time to tackle the challenges of semi-structured knowledge access. By continuing to develop and refine approaches like the one described in this post, we can unlock the vast amounts of knowledge contained in semi-structured documents and make them more accessible and useful for a wide range of applications.
