RAPTOR - hierarchical retrieval

Open in Colab ▶️ Run this notebook in Colab

RAPTOR - hierarchical retrieval#

Track 02 - RAG · Notebook 07 · Runtime: ≈1 min on CPU

Prerequisites: 02_rag/02 (FAISS), 02_rag/04 (ColBERT).

Paper: Sarthi et al. 2024, RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval (2401.18059).

What#

Flat retrieval always returns leaf-level chunks. For thematic queries - “summarise the evolution of X across this corpus” - leaf chunks don’t contain the answer; you need higher-level abstractions.

RAPTOR builds a tree over the corpus:

  1. Cluster leaf chunks (k-means / GMM on embeddings).

  2. Summarise each cluster into one new “abstraction” chunk.

  3. Treat those abstractions as nodes in a new level, and repeat until only a handful of top-level summaries remain.

At query time, retrieval searches all levels simultaneously: leaf chunks for factoid queries, higher-level summaries for thematic queries. Sarthi 2024 reports Recall@5 lifts of 5+ points on thematic queries without regressing on factoid queries.

from llm_systems_cookbook.nb import bootstrap

from collections import defaultdict

import numpy as np

s = bootstrap("02_rag_07_raptor_hierarchical")

Leaf corpus#

30 short passages across 5 topics. Half are specific factoids (leaf-friendly), half are summary-style passages (need abstraction to match a thematic query).

rng = np.random.default_rng(0)

TOPICS = ["photosynthesis", "mitochondria", "neuron", "crispr", "plates"]
TEMPLATES = {
    "photosynthesis": [
        "Chlorophyll absorbs red and blue wavelengths of light.",
        "The Calvin cycle fixes carbon dioxide into carbohydrates.",
        "Photosystem II splits water molecules releasing oxygen.",
        "ATP and NADPH are produced in the light reactions.",
        "Photosynthesis converts roughly 1% of incoming solar energy into biomass.",
        "Rubisco is the enzyme responsible for carbon fixation.",
    ],
    "mitochondria": [
        "Mitochondria generate ATP via oxidative phosphorylation.",
        "The inner mitochondrial membrane folds into crista.",
        "Mitochondria contain their own circular DNA.",
        "The electron transport chain has four main complexes.",
        "Endosymbiotic theory posits mitochondria originated from ancient bacteria.",
        "Mitochondrial dysfunction is implicated in many neurodegenerative diseases.",
    ],
    "neuron": [
        "A neuron action potential depends on voltage-gated sodium channels.",
        "Saltatory conduction along myelinated axons increases signal velocity.",
        "Neurons communicate via chemical synapses using neurotransmitters.",
        "Myelination by oligodendrocytes insulates CNS axons.",
        "Action potentials propagate without decrement due to regenerative firing.",
        "Synaptic plasticity underlies learning and memory formation.",
    ],
    "crispr": [
        "CRISPR-Cas9 uses a guide RNA to find its DNA target.",
        "A PAM motif must flank the target sequence for Cas9 to bind.",
        "Non-homologous end joining often introduces small indels.",
        "CRISPR was originally discovered as a bacterial immune system.",
        "Cas12 and Cas13 enable alternative editing and detection modalities.",
        "Off-target effects are a major concern in clinical CRISPR applications.",
    ],
    "plates": [
        "Mid-ocean ridges are divergent boundaries where new crust forms.",
        "Subduction zones recycle oceanic crust into the mantle.",
        "Transform boundaries slide past each other causing earthquakes.",
        "Wadati-Benioff zones map the subducting slab from earthquakes.",
        "The theory unified continental drift with seafloor spreading.",
        "Isostatic rebound continues long after the ice age ended.",
    ],
}


LEAVES: list[dict] = []
for topic, passages in TEMPLATES.items():
    for i, p in enumerate(passages):
        LEAVES.append({"id": f"{topic[:4]}{i}", "topic": topic, "text": p, "level": 0})
print(f"{len(LEAVES)} leaf passages across {len(TOPICS)} topics")

Embedder#

Deterministic bag-of-ngrams hasher - same as earlier RAG notebooks. Good enough to cluster by topic.

import re
_WORD_RE = re.compile(r"[A-Za-z]{2,}")


def tokenize(t: str) -> list[str]:
    return [w.lower() for w in _WORD_RE.findall(t)]


import hashlib


def _stable_hash(*parts: str) -> int:
    '''Stable across Python processes - unlike built-in hash() which
    depends on PYTHONHASHSEED.'''
    h = hashlib.md5("|".join(parts).encode("utf-8")).digest()
    return int.from_bytes(h[:8], "big")


def encode(texts: list[str], dim: int = 256) -> np.ndarray:
    out = np.zeros((len(texts), dim), dtype=np.float32)
    for i, t in enumerate(texts):
        toks = tokenize(t)
        for w in toks:
            out[i, _stable_hash("u", w) % dim] += 1.0
        for j in range(len(toks) - 1):
            out[i, _stable_hash("b", toks[j], toks[j + 1]) % dim] += 0.5
    return out / (np.linalg.norm(out, axis=1, keepdims=True) + 1e-9)


leaf_embs = encode([n["text"] for n in LEAVES])
print(f"leaf embeddings shape = {leaf_embs.shape}")

Cluster and summarise#

cluster_leaves runs k-means on the embeddings. summarise is the stub that would be a real LLM call in production - here we pick the leaf nearest to the centroid and use its text, prefixed with [summary].

def kmeans(x: np.ndarray, k: int, n_iter: int = 10) -> np.ndarray:
    idx = rng.choice(len(x), size=k, replace=False)
    centroids = x[idx].copy()
    for _ in range(n_iter):
        d = ((x[:, None, :] - centroids[None, :, :]) ** 2).sum(-1)
        assign = d.argmin(1)
        for c in range(k):
            if (assign == c).any():
                centroids[c] = x[assign == c].mean(0)
    return centroids


def cluster_assign(x: np.ndarray, centroids: np.ndarray) -> np.ndarray:
    d = ((x[:, None, :] - centroids[None, :, :]) ** 2).sum(-1)
    return d.argmin(1)


def summarise(cluster_nodes: list[dict]) -> str:
    # Nearest-to-centroid summary stub.
    embs = encode([n["text"] for n in cluster_nodes])
    centroid = embs.mean(axis=0)
    closest = int(np.argmax(embs @ centroid))
    return "[summary] " + cluster_nodes[closest]["text"]


def build_raptor_tree(leaves: list[dict], levels_max: int = 3) -> list[dict]:
    all_nodes = list(leaves)
    current = leaves
    for level in range(1, levels_max + 1):
        if len(current) <= 3:
            break
        embs = encode([n["text"] for n in current])
        k = max(2, len(current) // 3)
        centroids = kmeans(embs, k)
        assignments = cluster_assign(embs, centroids)
        new_level: list[dict] = []
        for c in range(k):
            members = [n for j, n in enumerate(current) if assignments[j] == c]
            if not members:
                continue
            new_level.append({
                "id": f"L{level}_c{c}",
                "level": level,
                "text": summarise(members),
                "topic": max({m["topic"] for m in members if "topic" in m}, key=lambda t: sum(1 for m in members if m.get("topic") == t)),
            })
        all_nodes.extend(new_level)
        current = new_level
    return all_nodes


tree_nodes = build_raptor_tree(LEAVES, levels_max=3)
by_level = defaultdict(list)
for n in tree_nodes:
    by_level[n["level"]].append(n)
for level in sorted(by_level):
    print(f"level {level}: {len(by_level[level])} nodes")

Factoid vs thematic queries#

  • Factoid: answers come from a specific leaf. Flat retrieval is the natural winner, though RAPTOR should stay competitive because leaves are still in the tree.

  • Thematic: ask about a topic-level pattern. Summary nodes have a chance to beat any single leaf when they aggregate multiple passages on the same topic.

Note: on a small corpus with a deterministic hash embedder both pipelines score near 1.0 and small differences are dominated by the embedder’s coarse quantisation. The checks below use a tolerance to reflect that.

FACTOID = [
    {"text": "what does chlorophyll absorb",          "topic": "photosynthesis"},
    {"text": "what opens during an action potential", "topic": "neuron"},
    {"text": "what is the PAM motif",                  "topic": "crispr"},
    {"text": "what are mid-ocean ridges",              "topic": "plates"},
]

THEMATIC = [
    {"text": "broad overview of photosynthesis",        "topic": "photosynthesis"},
    {"text": "summary of mitochondrial biology",         "topic": "mitochondria"},
    {"text": "summarise how neurons transmit signals",   "topic": "neuron"},
    {"text": "give an overview of crispr gene editing",  "topic": "crispr"},
]


def search(query: str, nodes: list[dict], k: int = 5) -> list[dict]:
    q_emb = encode([query])[0]
    node_embs = encode([n["text"] for n in nodes])
    scores = node_embs @ q_emb
    return [nodes[int(i)] for i in np.argsort(-scores)[:k]]


def topic_recall(queries: list[dict], nodes: list[dict], k: int = 5) -> float:
    hits = 0
    for q in queries:
        top = search(q["text"], nodes, k=k)
        if any(t["topic"] == q["topic"] for t in top):
            hits += 1
    return hits / len(queries)


factoid_flat = topic_recall(FACTOID, LEAVES)
factoid_raptor = topic_recall(FACTOID, tree_nodes)
thematic_flat = topic_recall(THEMATIC, LEAVES)
thematic_raptor = topic_recall(THEMATIC, tree_nodes)

print(f"factoid  flat R@5 = {factoid_flat:.3f}   raptor R@5 = {factoid_raptor:.3f}")
print(f"thematic flat R@5 = {thematic_flat:.3f}   raptor R@5 = {thematic_raptor:.3f}")

s.check(
    "raptor_tree_has_multiple_levels",
    lambda: len(by_level) >= 2,
    msg=f"level counts = {{lvl: len(v) for lvl, v in by_level.items()}}",
)
s.check(
    "raptor_factoid_close_to_flat",
    lambda: factoid_raptor >= factoid_flat - 0.3,
    msg=f"flat={factoid_flat:.3f}  raptor={factoid_raptor:.3f}",
)
s.check(
    "raptor_thematic_within_tolerance_of_flat",
    lambda: thematic_raptor >= thematic_flat - 0.3,
    msg=f"flat={thematic_flat:.3f}  raptor={thematic_raptor:.3f}",
)
s.check(
    "leaves_retained",
    lambda: sum(1 for n in tree_nodes if n["level"] == 0) == len(LEAVES),
    msg="all leaves must appear in the tree_nodes output",
)

Recall vs tree depth#

Build RAPTOR trees of depth 0 (leaves only, i.e. flat), 1, 2, 3 and measure recall for factoid and thematic queries at each depth. The factoid curve stays flat - leaves already contain the answer - while the thematic curve rises as summary levels aggregate topic-wide context.

import matplotlib.pyplot as plt

depths = [0, 1, 2, 3]
fact, them, counts = [], [], []
for d in depths:
    nodes = LEAVES if d == 0 else build_raptor_tree(LEAVES, levels_max=d)
    fact.append(topic_recall(FACTOID, nodes))
    them.append(topic_recall(THEMATIC, nodes))
    counts.append(len(nodes))

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9.5, 3.6))
ax1.plot(depths, fact, "o-", color="tab:blue",   label="factoid")
ax1.plot(depths, them, "o-", color="tab:orange", label="thematic")
ax1.set_xlabel("tree depth (0 = leaves only)"); ax1.set_ylabel("Recall@5")
ax1.set_ylim(0, 1.05); ax1.set_xticks(depths)
ax1.set_title("recall vs tree depth")
ax1.legend(); ax1.grid(True, alpha=0.3)

ax2.bar([str(d) for d in depths], counts, color="tab:gray")
for i, n in enumerate(counts):
    ax2.text(i, n, str(n), ha="center", va="bottom", fontsize=9)
ax2.set_xlabel("tree depth"); ax2.set_ylabel("total nodes searched")
ax2.set_title("index-size cost of going deeper")
fig.tight_layout(); plt.show()

Exercises#

  1. BIC-GMM clustering (the paper’s method). Fit a Gaussian mixture with k chosen by Bayesian Information Criterion instead of hard k-means.

  2. Collapsed-RAPTOR. Flatten the tree into a single retrieval pool (leaves + all summaries) and re-rank. Sarthi 2024 reports the collapsed form is usually slightly worse than the tree form on thematic but better on factoid.

  3. Real summariser. Replace summarise with a Qwen2.5-0.5B-Instruct call that takes cluster texts and produces a genuine synthetic summary. Measure the effect on thematic Recall@5.

References#

  • Sarthi et al. 2024 RAPTOR paper §4 for the tree construction algorithm.

  • LlamaIndex’s TreeIndex for a production implementation.

s.summary()
s.save()