(NOTE: if running on Colab, you will need to supply a WandB API Key in addition to your HFToken. Also, you'll need to change the runtime to a T4.)

Async Qdrant Knowledge Storage¶

Introduction¶

The fed-rag library supports a simple, in-memory knowledge store for rapid creation and development cycles of RAG systems. For larger scale fine-tuning jobs, you may need a more optimized knowledge store. FedRAG supports a seamless Qdrant integration in the form of the QdrantKnowledgeStore, allowing you to connect to any Qdrant service—whether running locally or in a managed/cloud environment.

In this notebook, we demonstrate how to launch a local Qdrant service and use it as the knowledge storage for your RAG system.

Install dependencies¶

The QdrantKnowledgeStore requires the installation of the qdrant extra. Note that we also will use a HuggingFace SentenceTransformer as the retriever/embedding model to encode our knowledge artifacts prior to loading them to our knowledge store.

In [ ]:

# If running in a Google Colab, the first attempt at installing fed-rag may fail,
# though for reasons unknown to me yet, if you try a second time, it magically works...
!uv pip install fed-rag[huggingface,qdrant] -q

# If running in a Google Colab, the first attempt at installing fed-rag may fail, # though for reasons unknown to me yet, if you try a second time, it magically works... !uv pip install fed-rag[huggingface,qdrant] -q

In [ ]:

# We'll use the docker SDK to launch the Qdrant docker image
!uv pip install docker -q

# We'll use the docker SDK to launch the Qdrant docker image !uv pip install docker -q

Launch a Local Qdrant Service (with Docker)¶

This step assumes that you have docker installed on your machine. If not installed, refer to the official Docker docs for installation found here.

IMPORTANT NOTE: if you are running this within a Google Colab, you won't be able to run a docker image. Instead, you can run the rest of this notebook by using an in-memory instance of Qdrant.

If using a Colab, set the WITH_DOCKER to False

In [1]:

WITH_DOCKER = True

WITH_DOCKER = True

In [2]:

if WITH_DOCKER:
    import docker
    import os
    import time

    client = docker.from_env()
    image_name = "qdrant/qdrant"

    # first see if we need to pull the docker image
    try:
        client.images.get(image_name)
        print(f"Image '{image_name}' already exists locally")
    except docker.errors.ImageNotFound:
        print(f"Image '{image_name}' not found locally. Pulling...")
        # Pull with progress information
        for line in client.api.pull(image_name, stream=True, decode=True):
            if "progress" in line:
                print(f"\r{line['status']}: {line['progress']}", end="")
            elif "status" in line:
                print(f"\r{line['status']}", end="")
        print("\nPull complete!")

    # run the Qdrant container
    container = client.containers.run(
        "qdrant/qdrant",
        detach=True,  # Run in background
        ports={"6333/tcp": 6333, "6334/tcp": 6334},
        volumes={
            f"{os.getcwd()}/qdrant_storage": {
                "bind": "/qdrant/storage",
                "mode": "rw",
            }
        },
        name="qdrant-demo-fedrag-nb",
    )

    print(f"Container started with ID: {container.id}")

    # wait a moment for the container to initialize
    time.sleep(3)

    # Check container status
    container.reload()  # Refresh container data
    print(f"Container status: {container.status}")
    print(f"Container logs:")
    print(container.logs().decode("utf-8"))

if WITH_DOCKER: import docker import os import time client = docker.from_env() image_name = "qdrant/qdrant" # first see if we need to pull the docker image try: client.images.get(image_name) print(f"Image '{image_name}' already exists locally") except docker.errors.ImageNotFound: print(f"Image '{image_name}' not found locally. Pulling...") # Pull with progress information for line in client.api.pull(image_name, stream=True, decode=True): if "progress" in line: print(f"\r{line['status']}: {line['progress']}", end="") elif "status" in line: print(f"\r{line['status']}", end="") print("\nPull complete!") # run the Qdrant container container = client.containers.run( "qdrant/qdrant", detach=True, # Run in background ports={"6333/tcp": 6333, "6334/tcp": 6334}, volumes={ f"{os.getcwd()}/qdrant_storage": { "bind": "/qdrant/storage", "mode": "rw", } }, name="qdrant-demo-fedrag-nb", ) print(f"Container started with ID: {container.id}") # wait a moment for the container to initialize time.sleep(3) # Check container status container.reload() # Refresh container data print(f"Container status: {container.status}") print(f"Container logs:") print(container.logs().decode("utf-8"))

Image 'qdrant/qdrant' already exists locally
Container started with ID: 7064b1dfed933eaefd7f49a04eca5c0f1785ce312c46d1f5c94adbd8259ef6e3
Container status: running
Container logs:
           _                 _    
  __ _  __| |_ __ __ _ _ __ | |_  
 / _` |/ _` | '__/ _` | '_ \| __| 
| (_| | (_| | | | (_| | | | | |_  
 \__, |\__,_|_|  \__,_|_| |_|\__| 
    |_|                           

Version: 1.14.0, build: 3617a011
Access web UI at http://localhost:6333/dashboard

2025-06-18T02:45:36.556498Z  INFO storage::content_manager::consensus::persistent: Loading raft state from ./storage/raft_state.json    
2025-06-18T02:45:36.559757Z  INFO qdrant: Distributed mode disabled    
2025-06-18T02:45:36.559783Z  INFO qdrant: Telemetry reporting enabled, id: 51912651-ded0-4255-b0f8-6c3506e2563c    
2025-06-18T02:45:36.559809Z  INFO qdrant: Inference service is not configured.    
2025-06-18T02:45:36.562017Z  INFO qdrant::actix: TLS disabled for REST API    
2025-06-18T02:45:36.562087Z  INFO qdrant::actix: Qdrant HTTP listening on 6333    
2025-06-18T02:45:36.562106Z  INFO actix_server::builder: Starting 11 workers
2025-06-18T02:45:36.562117Z  INFO actix_server::server: Actix runtime found; starting in Actix runtime
2025-06-18T02:45:36.566445Z  INFO qdrant::tonic: Qdrant gRPC listening on 6334    
2025-06-18T02:45:36.566456Z  INFO qdrant::tonic: TLS disabled for gRPC API    

Setup the Retriever and AsyncQdrantKnowledgeStore¶

In [3]:

from fed_rag.knowledge_stores import AsyncQdrantKnowledgeStore
from fed_rag.retrievers.huggingface import (
    HFSentenceTransformerRetriever,
)
from fed_rag.data_structures import KnowledgeNode, NodeType

from fed_rag.knowledge_stores import AsyncQdrantKnowledgeStore from fed_rag.retrievers.huggingface import ( HFSentenceTransformerRetriever, ) from fed_rag.data_structures import KnowledgeNode, NodeType

In [4]:

QUERY_ENCODER_NAME = "nthakur/dragon-plus-query-encoder"
CONTEXT_ENCODER_NAME = "nthakur/dragon-plus-context-encoder"

# retriever
retriever = HFSentenceTransformerRetriever(
    query_model_name=QUERY_ENCODER_NAME,
    context_model_name=CONTEXT_ENCODER_NAME,
    load_model_at_init=False,
)

# knowledge store
if WITH_DOCKER:
    knowledge_store = AsyncQdrantKnowledgeStore(
        collection_name="nthakur.dragon-plus-context-encoder"
    )
else:
    knowledge_store = AsyncQdrantKnowledgeStore(
        collection_name="nthakur.dragon-plus-context-encoder", in_memory=True
    )

QUERY_ENCODER_NAME = "nthakur/dragon-plus-query-encoder" CONTEXT_ENCODER_NAME = "nthakur/dragon-plus-context-encoder" # retriever retriever = HFSentenceTransformerRetriever( query_model_name=QUERY_ENCODER_NAME, context_model_name=CONTEXT_ENCODER_NAME, load_model_at_init=False, ) # knowledge store if WITH_DOCKER: knowledge_store = AsyncQdrantKnowledgeStore( collection_name="nthakur.dragon-plus-context-encoder" ) else: knowledge_store = AsyncQdrantKnowledgeStore( collection_name="nthakur.dragon-plus-context-encoder", in_memory=True )

Let's Add Some Knowledge¶

In [5]:

# a small sample from the Dec 2021 Wikipedia dump
text_chunks = [
    {
        "id": "140",
        "title": "History of marine biology",
        "section": "James Cook",
        "text": " James Cook is well known for his voyages of exploration for the British Navy in which he mapped out a significant amount of the world's uncharted waters. Cook's explorations took him around the world twice and led to countless descriptions of previously unknown plants and animals. Cook's explorations influenced many others and led to a number of scientists examining marine life more closely. Among those influenced was Charles Darwin who went on to make many contributions of his own. ",
    },
    {
        "id": "141",
        "title": "History of marine biology",
        "section": "Charles Darwin",
        "text": " Charles Darwin, best known for his theory of evolution, made many significant contributions to the early study of marine biology. He spent much of his time from 1831 to 1836 on the voyage of HMS Beagle collecting and studying specimens from a variety of marine organisms. It was also on this expedition where Darwin began to study coral reefs and their formation. He came up with the theory that the overall growth of corals is a balance between the growth of corals upward and the sinking of the sea floor. He then came up with the idea that wherever coral atolls would be found, the central island where the coral had started to grow would be gradually subsiding",
    },
    {
        "id": "142",
        "title": "History of marine biology",
        "section": "Charles Wyville Thomson",
        "text": " Another influential expedition was the voyage of HMS Challenger from 1872 to 1876, organized and later led by Charles Wyville Thomson. It was the first expedition purely devoted to marine science. The expedition collected and analyzed thousands of marine specimens, laying the foundation for present knowledge about life near the deep-sea floor. The findings from the expedition were a summary of the known natural, physical and chemical ocean science to that time.",
    },
]

# a small sample from the Dec 2021 Wikipedia dump text_chunks = [ { "id": "140", "title": "History of marine biology", "section": "James Cook", "text": " James Cook is well known for his voyages of exploration for the British Navy in which he mapped out a significant amount of the world's uncharted waters. Cook's explorations took him around the world twice and led to countless descriptions of previously unknown plants and animals. Cook's explorations influenced many others and led to a number of scientists examining marine life more closely. Among those influenced was Charles Darwin who went on to make many contributions of his own. ", }, { "id": "141", "title": "History of marine biology", "section": "Charles Darwin", "text": " Charles Darwin, best known for his theory of evolution, made many significant contributions to the early study of marine biology. He spent much of his time from 1831 to 1836 on the voyage of HMS Beagle collecting and studying specimens from a variety of marine organisms. It was also on this expedition where Darwin began to study coral reefs and their formation. He came up with the theory that the overall growth of corals is a balance between the growth of corals upward and the sinking of the sea floor. He then came up with the idea that wherever coral atolls would be found, the central island where the coral had started to grow would be gradually subsiding", }, { "id": "142", "title": "History of marine biology", "section": "Charles Wyville Thomson", "text": " Another influential expedition was the voyage of HMS Challenger from 1872 to 1876, organized and later led by Charles Wyville Thomson. It was the first expedition purely devoted to marine science. The expedition collected and analyzed thousands of marine specimens, laying the foundation for present knowledge about life near the deep-sea floor. The findings from the expedition were a summary of the known natural, physical and chemical ocean science to that time.", }, ]

In [6]:

from fed_rag.data_structures import KnowledgeNode, NodeType

# create knowledge nodes
nodes = []
texts = []
for c in text_chunks:
    text = c.pop("text")
    title = c.pop("title")
    section = c.pop("section")
    context_text = f"title: {title}\nsection: {section}\ntext: {text}"
    texts.append(context_text)

# batch encode
batch_embeddings = retriever.encode_context(texts)

for jx, c in enumerate(text_chunks):
    node = KnowledgeNode(
        embedding=batch_embeddings[jx].tolist(),
        node_type=NodeType.TEXT,
        text_content=texts[jx],
        metadata=c,
    )
    nodes.append(node)

from fed_rag.data_structures import KnowledgeNode, NodeType # create knowledge nodes nodes = [] texts = [] for c in text_chunks: text = c.pop("text") title = c.pop("title") section = c.pop("section") context_text = f"title: {title}\nsection: {section}\ntext: {text}" texts.append(context_text) # batch encode batch_embeddings = retriever.encode_context(texts) for jx, c in enumerate(text_chunks): node = KnowledgeNode( embedding=batch_embeddings[jx].tolist(), node_type=NodeType.TEXT, text_content=texts[jx], metadata=c, ) nodes.append(node)

In [7]:

await knowledge_store.load_nodes(nodes)

await knowledge_store.load_nodes(nodes)

In [8]:

# knowledge_store.count ## property `count` is not support for async version
await knowledge_store.get_count()

# knowledge_store.count ## property `count` is not support for async version await knowledge_store.get_count()

Out[8]:

3

Retriever From The Knowledge Store¶

In [9]:

query = "Who is James Cook?"
query_emb = retriever.encode_query(query).tolist()

query = "Who is James Cook?" query_emb = retriever.encode_query(query).tolist()

In [10]:

retrieved_nodes = await knowledge_store.retrieve(query_emb=query_emb, top_k=1)

retrieved_nodes = await knowledge_store.retrieve(query_emb=query_emb, top_k=1)

In [11]:

similarity_score, knowledge_node = retrieved_nodes[0]

similarity_score, knowledge_node = retrieved_nodes[0]

In [12]:

print("Similarity score: ", similarity_score)
print("KnowledgeNode: ", knowledge_node)

print("Similarity score: ", similarity_score) print("KnowledgeNode: ", knowledge_node)

Similarity score:  0.49984106
KnowledgeNode:  node_id='2f6f5007-7d97-474b-8612-2d7eade2a5ef' embedding=None node_type=<NodeType.TEXT: 'text'> text_content="title: History of marine biology\nsection: James Cook\ntext:  James Cook is well known for his voyages of exploration for the British Navy in which he mapped out a significant amount of the world's uncharted waters. Cook's explorations took him around the world twice and led to countless descriptions of previously unknown plants and animals. Cook's explorations influenced many others and led to a number of scientists examining marine life more closely. Among those influenced was Charles Darwin who went on to make many contributions of his own. " image_content=None metadata={'id': '140'}

Batch Retrieval¶

In [13]:

another_query = "Who is Charles Darwin?"
another_query_emb = retriever.encode_query(another_query).tolist()

another_query = "Who is Charles Darwin?" another_query_emb = retriever.encode_query(another_query).tolist()

In [14]:

batch_retrieval_result = await knowledge_store.batch_retrieve(
    query_embs=[query_emb, another_query_emb], top_k=1
)

batch_retrieval_result = await knowledge_store.batch_retrieve( query_embs=[query_emb, another_query_emb], top_k=1 )

In [15]:

# Who is James Cook?
batch_retrieval_result[0]

# Who is James Cook? batch_retrieval_result[0]

Out[15]:

[(0.49984106,
  KnowledgeNode(node_id='2f6f5007-7d97-474b-8612-2d7eade2a5ef', embedding=None, node_type=<NodeType.TEXT: 'text'>, text_content="title: History of marine biology\nsection: James Cook\ntext:  James Cook is well known for his voyages of exploration for the British Navy in which he mapped out a significant amount of the world's uncharted waters. Cook's explorations took him around the world twice and led to countless descriptions of previously unknown plants and animals. Cook's explorations influenced many others and led to a number of scientists examining marine life more closely. Among those influenced was Charles Darwin who went on to make many contributions of his own. ", image_content=None, metadata={'id': '140'}))]

In [16]:

# Who is Charles Darwin?
batch_retrieval_result[1]

# Who is Charles Darwin? batch_retrieval_result[1]

Out[16]:

[(0.50851303,
  KnowledgeNode(node_id='c1d34829-ff52-4fa5-95be-00088e7853e6', embedding=None, node_type=<NodeType.TEXT: 'text'>, text_content='title: History of marine biology\nsection: Charles Darwin\ntext:  Charles Darwin, best known for his theory of evolution, made many significant contributions to the early study of marine biology. He spent much of his time from 1831 to 1836 on the voyage of HMS Beagle collecting and studying specimens from a variety of marine organisms. It was also on this expedition where Darwin began to study coral reefs and their formation. He came up with the theory that the overall growth of corals is a balance between the growth of corals upward and the sinking of the sea floor. He then came up with the idea that wherever coral atolls would be found, the central island where the coral had started to grow would be gradually subsiding', image_content=None, metadata={'id': '141'}))]

Assembling an AsyncRAGSystem¶

Define an LLM Generator¶

In [17]:

from fed_rag.generators.huggingface import HFPretrainedModelGenerator
import torch
from transformers.generation.utils import GenerationConfig

generation_cfg = GenerationConfig(
    do_sample=True,
    eos_token_id=151643,
    bos_token_id=151643,
    max_new_tokens=2048,
    top_p=0.9,
    temperature=0.6,
    cache_implementation="offloaded",
    stop_strings="</response>",
)
generator = HFPretrainedModelGenerator(
    model_name="Qwen/Qwen2.5-0.5B",
    load_model_at_init=False,
    load_model_kwargs={"device_map": "auto", "torch_dtype": torch.float16},
    generation_config=generation_cfg,
)

from fed_rag.generators.huggingface import HFPretrainedModelGenerator import torch from transformers.generation.utils import GenerationConfig generation_cfg = GenerationConfig( do_sample=True, eos_token_id=151643, bos_token_id=151643, max_new_tokens=2048, top_p=0.9, temperature=0.6, cache_implementation="offloaded", stop_strings="", ) generator = HFPretrainedModelGenerator( model_name="Qwen/Qwen2.5-0.5B", load_model_at_init=False, load_model_kwargs={"device_map": "auto", "torch_dtype": torch.float16}, generation_config=generation_cfg, )

In [18]:

from fed_rag import AsyncRAGSystem, RAGConfig

rag_config = RAGConfig(top_k=2)
rag_system = AsyncRAGSystem(
    knowledge_store=knowledge_store,  # knowledge store loaded from knowledge_store.py
    generator=generator,
    retriever=retriever,
    rag_config=rag_config,
)

from fed_rag import AsyncRAGSystem, RAGConfig rag_config = RAGConfig(top_k=2) rag_system = AsyncRAGSystem( knowledge_store=knowledge_store, # knowledge store loaded from knowledge_store.py generator=generator, retriever=retriever, rag_config=rag_config, )

In [19]:

# test a query
response = await rag_system.query("Who is James Cook?")

# test a query response = await rag_system.query("Who is James Cook?")

The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
Setting `pad_token_id` to `eos_token_id`:151643 for open-end generation.
The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.

In [20]:

print(response)

print(response)

Assistant: James Cook was an explorer from the British Navy who made significant contributions to marine biology. He mapped out vast oceans and collected numerous marine specimens, laying the groundwork for present knowledge about marine life. His voyages influenced many others, including Charles Darwin, who made numerous contributions to marine science.

Clean up¶

In [21]:

await knowledge_store.clear()

await knowledge_store.clear()

In [22]:

if WITH_DOCKER:
    # stop and remove container
    container.stop()
    container.remove()

if WITH_DOCKER: # stop and remove container container.stop() container.remove()