Embeddings and Building a Database of Vectors

Milvus has a function called .from_documents. Basically, this takes documents, embeds them, and stores them into the Milvus database. We will use a basic structure to make sure all of our important information gets included.

In the code below, the mdw is my username. You can find your location by typing: echo $SCRATCH in the terminal.

# imports
from langchain_milvus import Milvus

vector_store = Milvus.from_documents(
    documents=all_splits,
    embedding=embed_model,
    collection_name=collection_name,
    connection_args={"uri":"/anvil/scratch/x-mdw/milvus_demo.db"},    # change this to your scratch directory
    drop_old=True,
)

The documents, in this case, would be the list of all of the chunks that resulted from splitting up your PDF. Each chunk then becomes one entry in the vector database.

For embedding, you should use the embed_model from Ollama, so that each of the chunks from all_splits gets converted into a vector.

We defined a collection name and URI way back at the start of these notes, when first directing the Milvus database to live in the SCRATCH directory.

Once vector_store has been created, we will need to turn the store into a retriever object. This will allow it to take a question, then find the most relevant chunks of the text. With this retriever, the next step is to create a retrieval chain, which will take the selected chunks, and use them to generate the response.

from langchain.chains import create_retrieval_chain

retriever = vector_store.as_retriever()
chain = create_retrieval_chain(combine_docs_chain=llm,retriever=retriever)

We will use an already made prompt template from LangChain Hub that is specifically designed for RAG. Its very base-level thinking is:

Use the following context to answer the question:

Context:
{context}

Question:
{input}

If the LLM had to process each retrieved chunk individually, the process could easily begin to take a long time. So instead, the context will consist of any and all selected relevant chunks.

from langchain import hub

retrieval_qa_chat_prompt = hub.pull("langchain-ai/retrieval-qa-chat")

(Dr Ward learned the method in the block before from Chat GPT!)

from langchain_core.runnables import RunnableLambda
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)
retrieval_chain = (
    {
        "context": lambda x: retriever.invoke(x["input"]),
        "input": lambda x: x["input"],
    }
    | RunnableLambda(
        lambda x: {
            "context": format_docs(x["context"]),
            "input": x["input"],
	}
    )
    | retrieval_qa_chat_prompt
    | llm
)

Now retrieval_chain will connect everything together! So the final pipeline should be:

User Question -> Retriever -> Combine Docs Chain -> LLM

To actually go and test the LLM:

# ask the model something about your selected PDF document
response = retrieval_chain.invoke({"input": "What is the main idea of this document?"})
print(response)