Build a RAG pipeline

Retrieval-augmented generation, or RAG, answers a question from your own documents instead of only the language model’s training data. Embed your documents with the GenAI API, store the embeddings in a TIR Vector Database, and pass the closest matches to a language model you run yourself.

Before you begin

You need:

• A Personal Access Token. See Get started with TIR GenAI API.
• A Qdrant database created in your project with TIR Vector Database, and its host and API key from the database’s Overview tab. See TIR Vector Database.
• Somewhere to run your language model, such as a TIR Instance with a GPU plan. The GenAI API supplies only the embeddings and the vector search; the language model that answers the question runs on your own compute.
• A Hugging Face account and access token, if the language model you choose needs one to download.
• The documents you want to answer questions from.

How a RAG pipeline fits together

Diagram as text

flowchart LR
    O["Your own code"]
    R["Retrieval: TIR Vector Database"]
    L["Language model: your own compute"]
    O <--> R
    O <--> L

Your own code ties the pipeline together. It takes a question and asks the vector database for the closest matching chunks of your documents. It then passes the question and those chunks to the language model, which writes the answer. The language model runs on your own compute, such as a TIR Instance.

You can write this code with a framework such as LangChain or Semantic Kernel, or in plain code of your own.

Turn your documents into a vector store

Diagram as text

flowchart LR
    A[Aggregate] --> C[Clean] --> Lo[Load] --> S[Split] --> E[Embed] --> St[Store]

Before you can answer questions from your documents, turn them into searchable vectors:

1. Aggregate the source documents you want the pipeline to answer from.
2. Clean each document: remove personally identifiable information, confidential information, and anything still in development.
3. Load the cleaned documents into memory. Common tools for this step include Unstructured, LlamaIndex, and LangChain’s own document loaders.
4. Split each document into chunks small enough to fit in a prompt.
5. Embed each chunk with the GenAI API’s embedding model.
6. Store each embedding in your TIR Vector Database. TIR’s Qdrant is one option; other vector stores in common use include Pinecone, Weaviate, FAISS, and Chroma.

Install the required packages

pip install -U bitsandbytes transformers peft accelerate datasets huggingface_hub
pip install -U langchain langchain-text-splitters qdrant-client
pip install e2enetworks

langchain-text-splitters is a separate package from langchain, so the second line installs both.

Log in to Hugging Face

If the language model you chose needs a Hugging Face account to download, log in with your Hugging Face access token before you load it:

from huggingface_hub import login

login(token="<your-hugging-face-token>")

This step uses your own Hugging Face account and token rather than your E2E credentials.

Create a collection for your embeddings

Before you run the pipeline, create a Qdrant collection to hold your document embeddings, the same way you would for a single embedding. For the steps, see Generate vector embeddings. The code that follows calls the collection example_collection.

Load the language model

Load the tokenizer and the model you chose, and build a text-generation pipeline from them:

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain_text_splitters import SentenceTransformersTokenTextSplitter

model_name = "<the language model you chose, for example mistralai/Mistral-7B-Instruct-v0.2>"

tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    add_bos_token=True,
    add_eos_token=True,
    trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    trust_remote_code=True,
)
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto",
)

text_splitter = SentenceTransformersTokenTextSplitter()

Connect to Qdrant and the GenAI API

from qdrant_client import QdrantClient
from e2enetworks.cloud import tir

qdrant_client = QdrantClient(
    host="<your-qdrant-host>",
    port=6333,
    api_key="<your-vector-database-api-key>",
)

tir.init(api_key="<your-api-key>", access_token="<your-auth-token>")
tir_client = tir.ModelAPIClient(project="<your-project-id>")

<your-qdrant-host> and <your-vector-database-api-key> are the host and API key from your Vector Database’s Overview tab. <your-auth-token> is the Auth Token from your Personal Access Token. Copy it from the Create Token dialog or, later, from the Personal Access Tokens list. <your-api-key> is that same token’s API Key. To get it, use the Download Config File action next to the token in the list; the file holds both values. <your-project-id> is the project the token belongs to.

Embed and store your documents

Define a function that embeds one piece of text through the GenAI API:

def get_vector(text):
    response = tir_client.infer(
        model_name="e5-mistral-7b-instruct",
        data={"prompt": text},
    )
    return response.outputs[0].data

Define a function that stores one embedding in Qdrant:

from qdrant_client.http.models import PointStruct
import time


def insert_vector(vector, text, chunk_id):
    point = PointStruct(
        id=int(time.time() * 1000) + chunk_id,
        vector=vector,
        payload={"data": text, "chunk_id": chunk_id},
    )
    qdrant_client.upsert(collection_name="example_collection", wait=False, points=[point])

Load each document, split it into chunks, embed each chunk, and store it:

import os

text_files_dir = "<path to the folder holding your cleaned, loaded documents>"

for filename in os.listdir(text_files_dir):
    filepath = os.path.join(text_files_dir, filename)
    with open(filepath, "r", encoding="utf-8") as file:
        text = file.read()
    if not text:
        continue
    chunks = text_splitter.split_text(text)
    for chunk_id, chunk in enumerate(chunks):
        vector = get_vector(chunk)
        if not vector:
            continue
        insert_vector(vector, chunk, chunk_id)

Keep the vector store up to date

If your source documents change regularly, re-run only the new or updated documents through this pipeline instead of re-embedding everything each time.

Answer a question from your documents

At question time, embed the question, search the vector store for the closest chunks, and pass them to the language model as context:

def answer_question(question):
    vector = get_vector(question)

    search_result = qdrant_client.query_points(
        collection_name="example_collection",
        query=vector,
        limit=2,
    )

    context = " ".join(point.payload["data"] for point in search_result.points)

    prompt = f"""
    <INST> {question}
    Context: {context}
    <SYS>Answer using only the context provided.</SYS></INST>
    """

    sequences = pipe(
        prompt,
        do_sample=True,
        max_new_tokens=1000,
        temperature=0.7,
        top_k=50,
        top_p=0.95,
        num_return_sequences=1,
        return_full_text=False,
    )
    print(sequences[0]["generated_text"])

The search returns an object whose points list holds the matches, each with a payload, a score, and an id; read the stored text from point.payload["data"].

RAG compared with fine-tuning

Fine-tuning trains a model further on your own data to improve it at a specific task. RAG instead gives an unmodified model up-to-date, specific context at the time you ask it a question, and lets you add, update, or remove that context while the model stays exactly as it is. You can also use both: a fine-tuned model can still be given RAG context.

Ask a question

Call the function with a question your documents can answer:

answer_question("<a question your documents can answer>")

The language model prints an answer built from the context your vector database returned. Ask a question unrelated to your documents to see how the model answers on its own.

To generate and store embeddings on their own, see Generate vector embeddings. For how the GenAI API’s calling methods differ, see Models and calling methods. If a call in this pipeline fails, see Troubleshooting TIR GenAI API.