This is a follow up of this article.

Overview of RAG pipelines

Most robust RAG pipelines in production have the following key components

Embedding strategies

Retrieval types and evaluation of context

Post Retrieval techniques

Quantitative metrics for user response

For complex use cases, where more than one type of data is being handled by the same RAG system , we can define separate workflows for each of the use case. An end-end robust RAG system looks like this-

Lets dive into each of the individual components of the workflow, i’ve already spoken about chunking strategy in the previous post in this series, the main focus of this is the Query Pre-processing ,post processing and workflow parts of RAG.

Query Preprocessing and Routing

There are different ways of pre-processing your queries such as

1. Query based Routing

2. Sub Query generation

3. HyDe transformation

4. Query Decomposition

There’s newer methods and techniques added into llama toolkit everyday. But for now lets deep dive into each of these techniques to see where what fits best for the use case.

Query Based Routing

Agent Assisted Workflows are very useful in cases where more than one type of information is embedded and a query might be required to be analyzed for intent before being part of a chain of different workflows . LlamaIndex has a Agent Query Pipeline, that helps define a DAG like workflow for query routing .

A working example of the notebook i’ve used to summarize a call transcript from Yahoo! Finance is here . However the parts of the notebook that are interesting are the way’s the agentic workflow gets stitched together

from llama_index.core.query_pipeline import QueryPipeline as QP

qp = QP(verbose=True,llm=OpenAI(model="gpt-3.5-turbo"))

qp.add_modules(
    {
        "agent_input": agent_input_component,
        "react_prompt": react_prompt_component,
        "llm":OpenAI(model="gpt-3.5-turbo"),
        "react_output_parser": parse_react_output,
        "run_tool": run_tool,
        "process_response": process_response,
        "process_agent_response": process_agent_response,
    }
)
qp.add_chain(["agent_input", "react_prompt", "llm", "react_output_parser"])

# add conditional link from react output to tool call (if not done)
qp.add_link(
    "react_output_parser",
    "run_tool",
    condition_fn=lambda x: not x["done"],
    input_fn=lambda x: x["reasoning_step"],
)
# add conditional link from react output to final response processing (if done)
qp.add_link(
    "react_output_parser",
    "process_response",
    condition_fn=lambda x: x["done"],
    input_fn=lambda x: x["reasoning_step"],
)

# whether response processing or tool output processing, add link to final agent response
qp.add_link("process_response", "process_agent_response")
qp.add_link("run_tool", "process_agent_response")

Each of the modules are then chained together in the form of a DAG as below

The user query then gets parsed to find the right “tool” or module defined in your DAG

If you are indexing and querying against multiple tenants (like in SaaS), and each tenant has their own workflow or integrations with other orchestrations e.g., a workflow on Uber data might require the summary to be emailed , where as a workflow on a health insurance- might require you to submit the user response to a different system as an example of where you might want to route based on the query .

Sub-Question generation

Information retrieval platforms based on multiple models perform with different levels of accuracy in IR tasks or NER tasks based on which LLM is being used. In the healthcare domain, when there is more than one information component requested in a user query, its been my experience that some of the SLM’s and older order LLMs (like 3.5) while are cost effective might require the user query to be split into sub questions. As an example.-

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.query_engine import SubQuestionQueryEngine
from llama_index.core.callbacks import CallbackManager, LlamaDebugHandler
from llama_index.core import Settings
llama_debug = LlamaDebugHandler(print_trace_on_end=True)
callback_manager = CallbackManager([llama_debug])

Settings.callback_manager = callback_manager
pg_essay = SimpleDirectoryReader(input_files=["/content/call_transcript_elv.pdf"]).load_data()

# build index and query engine
vector_query_engine = VectorStoreIndex.from_documents(
    pg_essay,
    use_async=True,
).as_query_engine()

# setup base query engine as tool
query_engine_tools = [
    QueryEngineTool(
        query_engine=vector_query_engine,
        metadata=ToolMetadata(
            name="Elevance finance call",
            description="Elevance earnings finance call transcript",
        ),
    ),
]

query_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=query_engine_tools,
    use_async=True,
)
response = query_engine.query(
    "What was gael and mark's overview of first quarter performance  "
)

The SubQuestionQueryEngine can now differentiate the two components of query as follows -

HyDE Transformation

HyDe works by asking an LLM at inference time to generate a hypothetical document that matches the user query.

Next, these hypothetical documents undergo a pivotal transformation through an encoder, which has been trained using an unsupervised contrastive learning method.

To further refine the search process, the average of the query’s embedding and the embeddings of the hypothetical documents is calculated. This resultant query embedding represents a distilled representation of the query’s intent, now imbued with the contextual nuances captured from the hypothetical documents.

The final step involves computing the similarity between the query embedding (composed of the query and hypothetical document embeddings) and the document embeddings derived from the corpus. By leveraging the same encoder for both hypothetical documents and the corpus, HyDE seamlessly maps the query to the most pertinent documents in the collection.

A sample working of hyde transformation is as follows, however a more detailed working sample including the transformation prompt is here .

elv_call_transcript = SimpleDirectoryReader(
        input_files=["/content/call_transcript_elv.pdf"]
    ).load_data()
transcript_index = VectorStoreIndex.from_documents(elv_call_transcript)
transcript_index.storage_context.persist(persist_dir="/content/")
transcript_engine = transcript_index.as_query_engine(similarity_top_k=3)
hyde = HyDEQueryTransform(include_original=True)
transcript_hyde_query_engine = TransformQueryEngine(transcript_engine, hyde)
query_engine_tools = [
    QueryEngineTool(
        query_engine=transcript_hyde_query_engine,
        metadata=ToolMetadata(
            name="call_transcript",
            description=(
                "Provides information about Elevance financials for year 2024. "
                "Use a detailed plain text question as input to the tool."
            ),
        ),
    )]

Query Decomposition

If sub question simply breaks down information from a user query, multi-step query transformation, is a more creative approach underpinning the self-ask method. This technique involves a language model autonomously posing and answering a series of follow-up questions before addressing the original query, enabling it to independently synthesize and integrate information.

Autonomous Inquiry

The self-ask method begins with the LLM generating additional, relevant questions based on the initial query.

Independent Learning and Synthesis

As the LLM answers these follow-up questions, it connects scattered facts, synthesizing insights along the way. This iterative process allows the LLM to unearth relationships and uncover insights that might remain hidden in a single-step approach.

Providing a Synthesized Answer

After this iterative exploration, the QueryEngine combines the independently gathered information to produce a more informed and cohesive response to the initial query.

Lets look at working example-

from llama_index.core.indices.query.query_transform.base import (
    StepDecomposeQueryTransform,
)
# gpt-3
step_decompose_transform_gpt3 = StepDecomposeQueryTransform(
    llm=gpt35, verbose=True
)
index_summary = "Used to answer questions about the document"
# set Logging to DEBUG for more detailed outputs
from llama_index.core.query_engine import MultiStepQueryEngine

query_engine = index.as_query_engine(llm=gpt35)
query_engine = MultiStepQueryEngine(
    query_engine=query_engine,
    query_transform=step_decompose_transform_gpt3,
    index_summary=index_summary,
)
response_gpt35 = query_engine.query(
    "what was gaels overview of the financial earnings",
)

The Response above is the query decomposition which is then again used by the LLM to better synthesize the answer as below

The complete working notebook is in the link here

Query Post Processing

In the post processing , there are two parts- Ranking results and evaluation. But since evaluation is a core area cutting across different components of RAG, the next part of the series will be a more comprehensive deep dive into evaluation metrics

Re-Ranking

While re-ranking has been hugely simplified by out of box customizations like Cohere Re-ranking and its integration with the llama ecosystem, its still a lot more fun (especially for trickier use cases) to train your cross encoder for re-ranking.

For the purposes of this experiment (also because im bound by org policy), ive used a toy dataset to demonstrate cross encoder training. The model used below is BERT ,but any of BERT based models like x-bert , roberta can be used in substitution.

"""
Methodology:
Three steps are followed for AugSBERT data-augmentation strategy with Domain Transfer / Cross-Domain -
1. Cross-Encoder aka BERT is trained over STSb (source) dataset.
2. Cross-Encoder is used to label QQP training (target) dataset (Assume no labels/no annotations are provided).
3. Bi-encoder aka SBERT is trained over the labeled QQP (target) dataset.

"""

from torch.utils.data import DataLoader
from sentence_transformers import models, losses, util, LoggingHandler, SentenceTransformer
from sentence_transformers.cross_encoder import CrossEncoder
from sentence_transformers.cross_encoder.evaluation import CECorrelationEvaluator
from sentence_transformers.evaluation import BinaryClassificationEvaluator
from sentence_transformers.readers import InputExample
from datetime import datetime
from zipfile import ZipFile
import logging
import csv
import sys
import torch
import math
import gzip
import os


# You can specify any huggingface/transformers pre-trained model here, for example, bert-base-uncased, roberta-base, xlm-roberta-base
model_name = "bert-base-uncased"
batch_size = 16
num_epochs = 1
max_seq_length = 128
use_cuda = torch.cuda.is_available()

###### Read Datasets ######
sts_dataset_path = "datasets/stsbenchmark.tsv.gz"
qqp_dataset_path = "quora-IR-dataset"


# Check if the STSb dataset exists. If not, download and extract it
if not os.path.exists(sts_dataset_path):
    util.http_get("https://sbert.net/datasets/stsbenchmark.tsv.gz", sts_dataset_path)

# Check if the QQP dataset exists. If not, download and extract
if not os.path.exists(qqp_dataset_path):
    logging.info("Dataset not found. Download")
    zip_save_path = "quora-IR-dataset.zip"
    util.http_get(url="https://sbert.net/datasets/quora-IR-dataset.zip", path=zip_save_path)
    with ZipFile(zip_save_path, "r") as zipIn:
        zipIn.extractall(qqp_dataset_path)

# Cross encoder will be saved in this location. Path will be like this output/cross-encoder/stsb_indomain_bert-base-uncased
# Knowing the path correctly will be helpful in next step
cross_encoder_path = (
    "output/cross-encoder/stsb_indomain_"
    + model_name.replace("/", "-")
)

# Bi encoder will be saved in this location. Path will be like this output/cross-encoder/qqp_cross_domain_bert-base-uncased
# Knowing the path correctly will be helpful in next step
bi_encoder_path = (
    "output/bi-encoder/qqp_cross_domain_"
    + model_name.replace("/", "-")
)

###### Cross-encoder (simpletransformers) ######

logging.info("Loading cross-encoder model: {}".format(model_name))
# Use Huggingface/transformers model (like BERT, RoBERTa, XLNet, XLM-R) for cross-encoder model
cross_encoder = CrossEncoder(model_name, num_labels=1)

###### Bi-encoder (sentence-transformers) ######

logging.info("Loading bi-encoder model: {}".format(model_name))

# Use Huggingface/transformers model (like BERT, RoBERTa, XLNet, XLM-R) for mapping tokens to embeddings
word_embedding_model = models.Transformer(model_name, max_seq_length=max_seq_length)

# Apply mean pooling to get one fixed sized sentence vector
pooling_model = models.Pooling(
    word_embedding_model.get_word_embedding_dimension(),
    pooling_mode_mean_tokens=True,
    pooling_mode_cls_token=False,
    pooling_mode_max_tokens=False,
)

bi_encoder = SentenceTransformer(modules=[word_embedding_model, pooling_model])


# Step 1: Train cross-encoder model with STSbenchmark

logging.info("Step 1: Train cross-encoder: {} with STSbenchmark (source dataset)".format(model_name))

gold_samples = []
dev_samples = []
test_samples = []

with gzip.open(sts_dataset_path, "rt", encoding="utf8") as fIn:
    reader = csv.DictReader(fIn, delimiter="\t", quoting=csv.QUOTE_NONE)
    for row in reader:
        score = float(row["score"]) / 5.0  # Normalize score to range 0 ... 1

        if row["split"] == "dev":
            dev_samples.append(InputExample(texts=[row["sentence1"], row["sentence2"]], label=score))
        elif row["split"] == "test":
            test_samples.append(InputExample(texts=[row["sentence1"], row["sentence2"]], label=score))
        else:
            # As we want to get symmetric scores, i.e. CrossEncoder(A,B) = CrossEncoder(B,A), we pass both combinations to the train set
            gold_samples.append(InputExample(texts=[row["sentence1"], row["sentence2"]], label=score))
            gold_samples.append(InputExample(texts=[row["sentence2"], row["sentence1"]], label=score))


# We wrap gold_samples (which is a List[InputExample]) into a pytorch DataLoader
train_dataloader = DataLoader(gold_samples, shuffle=True, batch_size=batch_size)


# We add an evaluator, which evaluates the performance during training
evaluator = CECorrelationEvaluator.from_input_examples(dev_samples, name="sts-dev")

# Configure the training
warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1)  # 10% of train data for warm-up
logging.info("Warmup-steps: {}".format(warmup_steps))

# Train the cross-encoder model
cross_encoder.fit(
    train_dataloader=train_dataloader,
    evaluator=evaluator,
    epochs=num_epochs,
    evaluation_steps=1000,
    warmup_steps=warmup_steps,
    output_path=cross_encoder_path,
)

# Step 2: Label QQP train dataset using cross-encoder (BERT) model

logging.info("Step 2: Label QQP (target dataset) with cross-encoder: {}".format(model_name))

cross_encoder = CrossEncoder(cross_encoder_path)

silver_data = []

with open(os.path.join(qqp_dataset_path, "classification/train_pairs.tsv"), encoding="utf8") as fIn:
    reader = csv.DictReader(fIn, delimiter="\t", quoting=csv.QUOTE_NONE)
    for row in reader:
        if row["is_duplicate"] == "1":
            silver_data.append([row["question1"], row["question2"]])

silver_scores = cross_encoder.predict(silver_data)

# All model predictions should be between [0,1]
assert all(0.0 <= score <= 1.0 for score in silver_scores)

binary_silver_scores = [1 if score >= 0.5 else 0 for score in silver_scores]

# Step 3: Train bi-encoder (SBERT) model with QQP dataset - Augmented SBERT

logging.info("Step 3: Train bi-encoder: {} over labeled QQP (target dataset)".format(model_name))

# Convert the dataset to a DataLoader ready for training
logging.info("Loading BERT labeled QQP dataset")
qqp_train_data = list(
    InputExample(texts=[data[0], data[1]], label=score) for (data, score) in zip(silver_data, binary_silver_scores)
)


train_dataloader = DataLoader(qqp_train_data, shuffle=True, batch_size=batch_size)
train_loss = losses.MultipleNegativesRankingLoss(bi_encoder)

###### Classification ######
# Given (quesiton1, question2), is this a duplicate or not?
# The evaluator will compute the embeddings for both questions and then compute
# a cosine similarity. If the similarity is above a threshold, we have a duplicate.
logging.info("Read QQP dev dataset")

dev_sentences1 = []
dev_sentences2 = []
dev_labels = []

with open(os.path.join(qqp_dataset_path, "classification/dev_pairs.tsv"), encoding="utf8") as fIn:
    reader = csv.DictReader(fIn, delimiter="\t", quoting=csv.QUOTE_NONE)
    for row in reader:
        dev_sentences1.append(row["question1"])
        dev_sentences2.append(row["question2"])
        dev_labels.append(int(row["is_duplicate"]))

evaluator = BinaryClassificationEvaluator(dev_sentences1, dev_sentences2, dev_labels)

# Configure the training.
warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1)  # 10% of train data for warm-up
logging.info("Warmup-steps: {}".format(warmup_steps))

# Train the bi-encoder model
bi_encoder.fit(
    train_objectives=[(train_dataloader, train_loss)],
    evaluator=evaluator,
    epochs=num_epochs,
    evaluation_steps=1000,
    warmup_steps=warmup_steps,
    output_path=bi_encoder_path,
)

# Evaluate Augmented SBERT performance on QQP benchmark dataset

# Loading the augmented sbert model
bi_encoder = SentenceTransformer(bi_encoder_path)

logging.info("Read QQP test dataset")
test_sentences1 = []
test_sentences2 = []
test_labels = []

with open(os.path.join(qqp_dataset_path, "classification/test_pairs.tsv"), encoding="utf8") as fIn:
    reader = csv.DictReader(fIn, delimiter="\t", quoting=csv.QUOTE_NONE)
    for row in reader:
        test_sentences1.append(row["question1"])
        test_sentences2.append(row["question2"])
        test_labels.append(int(row["is_duplicate"]))

evaluator = BinaryClassificationEvaluator(test_sentences1, test_sentences2, test_labels)
bi_encoder.evaluate(evaluator)

Once the cross encoder model is trained, you can use this in conjunction with llamaindex as below

import os

from llama_index.core import VectorStoreIndex,SimpleDirectoryReader
from llama_index.core.postprocessor import SentenceTransformerRerank
import openai

openai.api_key = "sk..."
# Load documents from a dtputirectory
documents = SimpleDirectoryReader('datasets/data').load_data()

#Now lets use the bi encoder trained ealier
print("######## Using bi-encoder trained with our data")

index = VectorStoreIndex(
        documents, embed_model="local:output/bi-encoder/qqp_cross_domain_bert-base-uncased", show_progress=True
    )
query_engine = index.as_query_engine()

response = query_engine.query("What did the author do growing up?")
print(response)

#Now lets use the cross encoder trained ealier for re-ranking
print("######## Using cross-encoder trained with our data for reranking")

# Initialize the reranker
rerank = SentenceTransformerRerank(
    model="./output/cross-encoder/stsb_indomain_bert-base-uncased", top_n=3
)


# Build the query engine
query_engine = index.as_query_engine(similarity_top_k=10, node_postprocessors=[rerank])

response = query_engine.query("What did the author do growing up?")
print(response)

     

While these are good basics to get started, the most pertinent of them all is the evaluation(which i think deserves its own series ) but will be in a follow up to this three part write-up.

References

• Custom Cohere Reranker - LlamaIndex
• Precise Zero-Shot Dense Retrieval without Relevance Labels

https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/tools/llama-index-tools-google/examples/advanced_tools_usage.ipynb

With the increase in the number of LLM’s released per day, and the RAG approaches, I’m excited to explore an end-end approach on decision making when it comes to RAGs

Key Drivers

We’ll look into some of they key considerations that go into deciding which approach to take — this is assuming your organization has hosted instances/access to LLM’s and VDB’s

Cost

Data processing pipelines

Selecting a RAG approach

Evaluation metrics

Cost

Cost of LLM’s is one of the main key considerations into maintaining your RAG pipeline or even beginning to consider a use case for automation. How do you factor in cost/token into your RAG pipeline? There are out-of-the box token solutions like docsbot that have an up to date $/token

But this is a first step into understanding how much your model is going to cost per query . When migrating and switching to more cost effective models, the prompts can get quite verbose, if you are choosing to use this on specific domain use cases . eg., while gpt-4 might be able to understand context from a medical domain, gpt-3.5 might require a few shot examples to understand .

overall cost =(No.of documents in a day)*(no of queries/document)*(no of tokens/query)*(cost_per_token)

The number of tokens depends on the context-length of each model. packages like tiktoken offer out of the box calculations for number of tokens consumed/query.

Data Ingestion Pipelines

Data pre-processing and indexing are key to building effective (Retrieval Augmented Generation (RAG) pipelines. Before we go into an in-depth view of RAG, it is important to invest effort in building efficient data processing steps that will help in extracting the right data before we pass it to LLM to better utilize the information .

There are many libraries like unstructured and NLMIngestor that help extract sections, tables, images from pdfs. A Naive RAG pipeline would probably chunk and index a document and then retrieve a context based on a query.

Load from PDF

Note that the NLMIngestor (which is used by llmsherpa below) is Apache2 licensed and can be hosted in your own environment

llmsherpa_api_url = "https://readers.llmsherpa.com/api/document/developer/parseDocument?renderFormat=all"
pdf_url = "https://arxiv.org/pdf/1910.13461.pdf" # also allowed is a file path e.g. /home/downloads/xyz.pdf
pdf_reader = LayoutPDFReader(llmsherpa_api_url)
doc = pdf_reader.read_pdf(pdf_url)

Index and Query

The example used Qdrant in memory but the configuration can be changed when using in production

qdrant = QdrantClient(":memory:") # Create in-memory Qdrant instance, for testing, CI/CD
llmsherpa_api_url = "https://readers.llmsherpa.com/api/document/developer/parseDocument?renderFormat=all"
pdf_url = "https://arxiv.org/pdf/2102.12060.pdf" # also allowed is a file path e.g. /home/downloads/xyz.pdf
pdf_reader = LayoutPDFReader(llmsherpa_api_url)
doc = pdf_reader.read_pdf(pdf_url)
metadata,documents,idx=[],[],[]
for section in doc.sections():
  if section.children:
     metadata.append({"heading":section.to_text(include_children=False)})
     print(f"The heading is {section.to_text(include_children=False)}")
     documents.append(''.join(str(chunk.to_text()) for chunk in section.chunks()))
     idx.append(section.page_idx)

qdrant.add(collection_name="test_collection",documents=documents,metadata=metadata,ids=idx)

The end -end flow to index and query is in colab-notebook.

RAG and others..

The other part of this series will focus on RAG , different types of RAG and the Evaluation metrics.