When doing small-to-big retrieval, how

bbixqu

When doing small-to-big retrieval, how do I store my vectors in a custom vector database?

If I do regular RAG (top-k, similar sized chunks) I simply write:

Plain Text

VectorStoreIndex.from_documents(documents,
                                        storage_context=self.storage_context,
                                        service_context=service_context,
                                        show_progress=False)

where I provide my storage context directing me to the custom vector DB I am running locally (PGVector).

3 comments

bbixqu

Now I want to do small-to-big RAG. My function so far is this:

Plain Text

def small_to_big(self, query, documents, rag=False, chunk_size=1024, tree_depth=2, similarity_top_k=3):
        node_parser = SimpleNodeParser.from_defaults(chunk_size=chunk_size)
        base_nodes = node_parser.get_nodes_from_documents(documents)
        # Assign unique ids to each node
        for idx, node in enumerate(base_nodes):
            node.id_ = f"node-{idx}"
        prompt_helper = PromptHelper(context_window=4096, 
                                          num_output=256, 
                                          chunk_overlap_ratio=0.1, 
                                          chunk_size_limit=None)

        service_context = ServiceContext.from_defaults(llm=self.llm,
                                                        embed_model=self.embed_model, 
                                                        prompt_helper=prompt_helper)

        sub_chunk_sizes = [chunk_size // 2**i for i in range(1,tree_depth+1)]

        sub_node_parsers = [SimpleNodeParser.from_defaults(chunk_size=c) for c in sub_chunk_sizes]

        all_nodes = []
        for base_node in base_nodes:
            for n in sub_node_parsers:
                sub_nodes = n.get_nodes_from_documents([base_node])
                sub_inodes = [IndexNode.from_text_node(sn, base_node.node_id) for sn in sub_nodes]
                all_nodes.extend(sub_inodes)

            original_node = IndexNode.from_text_node(base_node, base_node.node_id)
            all_nodes.append(original_node)
        all_nodes_dict = {n.node_id: n for n in all_nodes}

        vector_index_chunk = VectorStoreIndex(all_nodes, 
                                              service_context=service_context)
        # vector_index_chunk = VectorStoreIndex(all_nodes, 
        #                                       service_context=service_context,
        #                                       storage_context=self.storage_context)

bbixqu

Plain Text

        
        vector_retriever_chunk = vector_index_chunk.as_retriever(similarity_top_k=similarity_top_k)
        retriever_chunk = RecursiveRetriever("vector",
                                             retriever_dict={"vector": vector_retriever_chunk},
                                             node_dict=all_nodes_dict,
                                             verbose=True)

        if rag:
            query_engine_chunk = RetrieverQueryEngine.from_args(retriever_chunk, 
                                                                service_context=service_context)
            response = query_engine_chunk.query(query)
            return str(response)

This works , but does not use my local vector store at all. I tried passing storage_context to VectorStoreIndex (the commented part) but then it does not work, it returns that the node indices were not found. How can I modify this function to work with my local vector db?

bbixqu

I see in the documentation

https://docs.llamaindex.ai/en/stable/module_guides/indexing/vector_store_index.html#using-vectorstoreindex

this example:

Plain Text

from llama_index.schema import TextNode

node1 = TextNode(text="<text_chunk>", id_="<node_id>")
node2 = TextNode(text="<text_chunk>", id_="<node_id>")
nodes = [node1, node2]
index = VectorStoreIndex(nodes)

can storage_context be passed to this to store the nodes locally in my vector db?

Add a reply

Find answers from the community

When doing small-to-big retrieval, how