Lesson 4: Auto-merging Retrieval¶

In [ ]:

import warnings
warnings.filterwarnings('ignore')

import warnings warnings.filterwarnings('ignore')

In [ ]:

import utils

import os
import openai
openai.api_key = utils.get_openai_api_key()

import utils import os import openai openai.api_key = utils.get_openai_api_key()

In [ ]:

from llama_index import SimpleDirectoryReader

documents = SimpleDirectoryReader(
    input_files=["./eBook-How-to-Build-a-Career-in-AI.pdf"]
).load_data()

from llama_index import SimpleDirectoryReader documents = SimpleDirectoryReader( input_files=["./eBook-How-to-Build-a-Career-in-AI.pdf"] ).load_data()

In [ ]:

print(type(documents), "\n")
print(len(documents), "\n")
print(type(documents[0]))
print(documents[0])

print(type(documents), "\n") print(len(documents), "\n") print(type(documents[0])) print(documents[0])

Auto-merging retrieval setup¶

In [ ]:

from llama_index import Document

document = Document(text="\n\n".join([doc.text for doc in documents]))

from llama_index import Document document = Document(text="\n\n".join([doc.text for doc in documents]))

In [ ]:

from llama_index.node_parser import HierarchicalNodeParser

# create the hierarchical node parser w/ default settings
node_parser = HierarchicalNodeParser.from_defaults(
    chunk_sizes=[2048, 512, 128]
)

from llama_index.node_parser import HierarchicalNodeParser # create the hierarchical node parser w/ default settings node_parser = HierarchicalNodeParser.from_defaults( chunk_sizes=[2048, 512, 128] )

In [ ]:

nodes = node_parser.get_nodes_from_documents([document])

nodes = node_parser.get_nodes_from_documents([document])

In [ ]:

from llama_index.node_parser import get_leaf_nodes

leaf_nodes = get_leaf_nodes(nodes)
print(leaf_nodes[30].text)

from llama_index.node_parser import get_leaf_nodes leaf_nodes = get_leaf_nodes(nodes) print(leaf_nodes[30].text)

In [ ]:

nodes_by_id = {node.node_id: node for node in nodes}

parent_node = nodes_by_id[leaf_nodes[30].parent_node.node_id]
print(parent_node.text)

nodes_by_id = {node.node_id: node for node in nodes} parent_node = nodes_by_id[leaf_nodes[30].parent_node.node_id] print(parent_node.text)

Building the index¶

In [ ]:

from llama_index.llms import OpenAI

llm = OpenAI(model="gpt-3.5-turbo", temperature=0.1)

from llama_index.llms import OpenAI llm = OpenAI(model="gpt-3.5-turbo", temperature=0.1)

In [ ]:

from llama_index import ServiceContext

auto_merging_context = ServiceContext.from_defaults(
    llm=llm,
    embed_model="local:BAAI/bge-small-en-v1.5",
    node_parser=node_parser,
)

from llama_index import ServiceContext auto_merging_context = ServiceContext.from_defaults( llm=llm, embed_model="local:BAAI/bge-small-en-v1.5", node_parser=node_parser, )

In [ ]:

from llama_index import VectorStoreIndex, StorageContext

storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)

automerging_index = VectorStoreIndex(
    leaf_nodes, storage_context=storage_context, service_context=auto_merging_context
)

automerging_index.storage_context.persist(persist_dir="./merging_index")

from llama_index import VectorStoreIndex, StorageContext storage_context = StorageContext.from_defaults() storage_context.docstore.add_documents(nodes) automerging_index = VectorStoreIndex( leaf_nodes, storage_context=storage_context, service_context=auto_merging_context ) automerging_index.storage_context.persist(persist_dir="./merging_index")

In [ ]:

# This block of code is optional to check
# if an index file exist, then it will load it
# if not, it will rebuild it

import os
from llama_index import VectorStoreIndex, StorageContext, load_index_from_storage
from llama_index import load_index_from_storage

if not os.path.exists("./merging_index"):
    storage_context = StorageContext.from_defaults()
    storage_context.docstore.add_documents(nodes)

    automerging_index = VectorStoreIndex(
            leaf_nodes,
            storage_context=storage_context,
            service_context=auto_merging_context
        )

    automerging_index.storage_context.persist(persist_dir="./merging_index")
else:
    automerging_index = load_index_from_storage(
        StorageContext.from_defaults(persist_dir="./merging_index"),
        service_context=auto_merging_context
    )

# This block of code is optional to check # if an index file exist, then it will load it # if not, it will rebuild it import os from llama_index import VectorStoreIndex, StorageContext, load_index_from_storage from llama_index import load_index_from_storage if not os.path.exists("./merging_index"): storage_context = StorageContext.from_defaults() storage_context.docstore.add_documents(nodes) automerging_index = VectorStoreIndex( leaf_nodes, storage_context=storage_context, service_context=auto_merging_context ) automerging_index.storage_context.persist(persist_dir="./merging_index") else: automerging_index = load_index_from_storage( StorageContext.from_defaults(persist_dir="./merging_index"), service_context=auto_merging_context )

Defining the retriever and running the query engine¶

In [ ]:

from llama_index.indices.postprocessor import SentenceTransformerRerank
from llama_index.retrievers import AutoMergingRetriever
from llama_index.query_engine import RetrieverQueryEngine

automerging_retriever = automerging_index.as_retriever(
    similarity_top_k=12
)

retriever = AutoMergingRetriever(
    automerging_retriever, 
    automerging_index.storage_context, 
    verbose=True
)

rerank = SentenceTransformerRerank(top_n=6, model="BAAI/bge-reranker-base")

auto_merging_engine = RetrieverQueryEngine.from_args(
    automerging_retriever, node_postprocessors=[rerank]
)

from llama_index.indices.postprocessor import SentenceTransformerRerank from llama_index.retrievers import AutoMergingRetriever from llama_index.query_engine import RetrieverQueryEngine automerging_retriever = automerging_index.as_retriever( similarity_top_k=12 ) retriever = AutoMergingRetriever( automerging_retriever, automerging_index.storage_context, verbose=True ) rerank = SentenceTransformerRerank(top_n=6, model="BAAI/bge-reranker-base") auto_merging_engine = RetrieverQueryEngine.from_args( automerging_retriever, node_postprocessors=[rerank] )

In [ ]:

auto_merging_response = auto_merging_engine.query(
    "What is the importance of networking in AI?"
)

auto_merging_response = auto_merging_engine.query( "What is the importance of networking in AI?" )

In [ ]:

from llama_index.response.notebook_utils import display_response

display_response(auto_merging_response)

from llama_index.response.notebook_utils import display_response display_response(auto_merging_response)

Putting it all Together¶

In [ ]:

import os

from llama_index import (
    ServiceContext,
    StorageContext,
    VectorStoreIndex,
    load_index_from_storage,
)
from llama_index.node_parser import HierarchicalNodeParser
from llama_index.node_parser import get_leaf_nodes
from llama_index import StorageContext, load_index_from_storage
from llama_index.retrievers import AutoMergingRetriever
from llama_index.indices.postprocessor import SentenceTransformerRerank
from llama_index.query_engine import RetrieverQueryEngine


def build_automerging_index(
    documents,
    llm,
    embed_model="local:BAAI/bge-small-en-v1.5",
    save_dir="merging_index",
    chunk_sizes=None,
):
    chunk_sizes = chunk_sizes or [2048, 512, 128]
    node_parser = HierarchicalNodeParser.from_defaults(chunk_sizes=chunk_sizes)
    nodes = node_parser.get_nodes_from_documents(documents)
    leaf_nodes = get_leaf_nodes(nodes)
    merging_context = ServiceContext.from_defaults(
        llm=llm,
        embed_model=embed_model,
    )
    storage_context = StorageContext.from_defaults()
    storage_context.docstore.add_documents(nodes)

    if not os.path.exists(save_dir):
        automerging_index = VectorStoreIndex(
            leaf_nodes, storage_context=storage_context, service_context=merging_context
        )
        automerging_index.storage_context.persist(persist_dir=save_dir)
    else:
        automerging_index = load_index_from_storage(
            StorageContext.from_defaults(persist_dir=save_dir),
            service_context=merging_context,
        )
    return automerging_index


def get_automerging_query_engine(
    automerging_index,
    similarity_top_k=12,
    rerank_top_n=6,
):
    base_retriever = automerging_index.as_retriever(similarity_top_k=similarity_top_k)
    retriever = AutoMergingRetriever(
        base_retriever, automerging_index.storage_context, verbose=True
    )
    rerank = SentenceTransformerRerank(
        top_n=rerank_top_n, model="BAAI/bge-reranker-base"
    )
    auto_merging_engine = RetrieverQueryEngine.from_args(
        retriever, node_postprocessors=[rerank]
    )
    return auto_merging_engine

import os from llama_index import ( ServiceContext, StorageContext, VectorStoreIndex, load_index_from_storage, ) from llama_index.node_parser import HierarchicalNodeParser from llama_index.node_parser import get_leaf_nodes from llama_index import StorageContext, load_index_from_storage from llama_index.retrievers import AutoMergingRetriever from llama_index.indices.postprocessor import SentenceTransformerRerank from llama_index.query_engine import RetrieverQueryEngine def build_automerging_index( documents, llm, embed_model="local:BAAI/bge-small-en-v1.5", save_dir="merging_index", chunk_sizes=None, ): chunk_sizes = chunk_sizes or [2048, 512, 128] node_parser = HierarchicalNodeParser.from_defaults(chunk_sizes=chunk_sizes) nodes = node_parser.get_nodes_from_documents(documents) leaf_nodes = get_leaf_nodes(nodes) merging_context = ServiceContext.from_defaults( llm=llm, embed_model=embed_model, ) storage_context = StorageContext.from_defaults() storage_context.docstore.add_documents(nodes) if not os.path.exists(save_dir): automerging_index = VectorStoreIndex( leaf_nodes, storage_context=storage_context, service_context=merging_context ) automerging_index.storage_context.persist(persist_dir=save_dir) else: automerging_index = load_index_from_storage( StorageContext.from_defaults(persist_dir=save_dir), service_context=merging_context, ) return automerging_index def get_automerging_query_engine( automerging_index, similarity_top_k=12, rerank_top_n=6, ): base_retriever = automerging_index.as_retriever(similarity_top_k=similarity_top_k) retriever = AutoMergingRetriever( base_retriever, automerging_index.storage_context, verbose=True ) rerank = SentenceTransformerRerank( top_n=rerank_top_n, model="BAAI/bge-reranker-base" ) auto_merging_engine = RetrieverQueryEngine.from_args( retriever, node_postprocessors=[rerank] ) return auto_merging_engine

In [ ]:

from llama_index.llms import OpenAI

index = build_automerging_index(
    [document],
    llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
    save_dir="./merging_index",
)

from llama_index.llms import OpenAI index = build_automerging_index( [document], llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1), save_dir="./merging_index", )

In [ ]:

query_engine = get_automerging_query_engine(index, similarity_top_k=6)

query_engine = get_automerging_query_engine(index, similarity_top_k=6)

TruLens Evaluation¶

In [ ]:

from trulens_eval import Tru

Tru().reset_database()

from trulens_eval import Tru Tru().reset_database()

Two layers¶

In [ ]:

auto_merging_index_0 = build_automerging_index(
    documents,
    llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
    embed_model="local:BAAI/bge-small-en-v1.5",
    save_dir="merging_index_0",
    chunk_sizes=[2048,512],
)

auto_merging_index_0 = build_automerging_index( documents, llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1), embed_model="local:BAAI/bge-small-en-v1.5", save_dir="merging_index_0", chunk_sizes=[2048,512], )

In [ ]:

auto_merging_engine_0 = get_automerging_query_engine(
    auto_merging_index_0,
    similarity_top_k=12,
    rerank_top_n=6,
)

auto_merging_engine_0 = get_automerging_query_engine( auto_merging_index_0, similarity_top_k=12, rerank_top_n=6, )

In [ ]:

from utils import get_prebuilt_trulens_recorder

tru_recorder = get_prebuilt_trulens_recorder(
    auto_merging_engine_0,
    app_id ='app_0'
)

from utils import get_prebuilt_trulens_recorder tru_recorder = get_prebuilt_trulens_recorder( auto_merging_engine_0, app_id ='app_0' )

In [ ]:

eval_questions = []
with open('generated_questions.text', 'r') as file:
    for line in file:
        # Remove newline character and convert to integer
        item = line.strip()
        eval_questions.append(item)

eval_questions = [] with open('generated_questions.text', 'r') as file: for line in file: # Remove newline character and convert to integer item = line.strip() eval_questions.append(item)

In [ ]:

def run_evals(eval_questions, tru_recorder, query_engine):
    for question in eval_questions:
        with tru_recorder as recording:
            response = query_engine.query(question)

def run_evals(eval_questions, tru_recorder, query_engine): for question in eval_questions: with tru_recorder as recording: response = query_engine.query(question)

In [ ]:

run_evals(eval_questions, tru_recorder, auto_merging_engine_0)

run_evals(eval_questions, tru_recorder, auto_merging_engine_0)

In [ ]:

from trulens_eval import Tru

Tru().get_leaderboard(app_ids=[])

from trulens_eval import Tru Tru().get_leaderboard(app_ids=[])

In [ ]:

Tru().run_dashboard()

Tru().run_dashboard()

Three layers¶

In [ ]:

auto_merging_index_1 = build_automerging_index(
    documents,
    llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
    embed_model="local:BAAI/bge-small-en-v1.5",
    save_dir="merging_index_1",
    chunk_sizes=[2048,512,128],
)

auto_merging_index_1 = build_automerging_index( documents, llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1), embed_model="local:BAAI/bge-small-en-v1.5", save_dir="merging_index_1", chunk_sizes=[2048,512,128], )

In [ ]:

auto_merging_engine_1 = get_automerging_query_engine(
    auto_merging_index_1,
    similarity_top_k=12,
    rerank_top_n=6,
)

auto_merging_engine_1 = get_automerging_query_engine( auto_merging_index_1, similarity_top_k=12, rerank_top_n=6, )

In [ ]:

tru_recorder = get_prebuilt_trulens_recorder(
    auto_merging_engine_1,
    app_id ='app_1'
)

tru_recorder = get_prebuilt_trulens_recorder( auto_merging_engine_1, app_id ='app_1' )

In [ ]:

run_evals(eval_questions, tru_recorder, auto_merging_engine_1)

run_evals(eval_questions, tru_recorder, auto_merging_engine_1)

In [ ]:

from trulens_eval import Tru

Tru().get_leaderboard(app_ids=[])

from trulens_eval import Tru Tru().get_leaderboard(app_ids=[])

In [ ]:

Tru().run_dashboard()

Tru().run_dashboard()