Treehop Rag

🚀 TreeHop：高效生成和過濾多跳問答的下一個查詢嵌入

TreeHop 是一個輕量級的嵌入級框架，旨在解決檢索增強生成（RAG）領域中傳統遞歸檢索範式的計算效率低下問題。通過消除基於大語言模型（LLM）的迭代查詢重寫需求，TreeHop 在保持先進性能的同時顯著降低了延遲。它通過動態查詢嵌入更新和剪枝策略，實現了簡化的“檢索 - 嵌入 - 檢索”工作流程。

項目鏈接

• GitHub
• arXiv
• HuggingFace

🚀 快速開始

系統要求

• 操作系統：Ubuntu 18.06 LTS 及以上版本，或 MacOS Big Sur 及以上版本。
• 硬件要求：Nvidia GPU 或配備 32GB 內存的 Apple Metal，系統內存 16GB 用於復現論文，64GB 用於訓練 TreeHop，硬盤需有 50GB 可用空間。
• Python 環境：Python 3.9 及以上版本，具體依賴請參考 requirements.txt。

準備工作

本倉庫包含用於復現的評估嵌入數據庫。激活 git LFS 後，使用 git lfs clone [倉庫鏈接] 克隆倉庫，或在現有本地倉庫中使用以下命令拉取數據：

git lfs pull

多跳檢索使用示例

以下示例使用 MultiHop RAG 評估數據集。倉庫中包含運行示例所需的文件，請參考準備工作。

from tree_hop import TreeHopModel
from passage_retrieval import MultiHopRetriever

EVALUATE_DATASET = "multihop_rag"

# 從 HuggingFace 加載 TreeHop 模型
tree_hop_model = TreeHopModel.from_pretrained("allen-li1231/treehop-rag")

# 加載檢索器
retriever = MultiHopRetriever(
    "BAAI/bge-m3",
    passages=f"embedding_data/{EVALUATE_DATASET}/eval_passages.jsonl",
    passage_embeddings=f"embedding_data/{EVALUATE_DATASET}/eval_content_dense.npy",
    # 如果 faiss 索引已初始化，請取消註釋以下行以加快加載速度
    # faiss_index=f"embedding_data/{EVALUATE_DATASET}/index.faiss",
    tree_hop_model=tree_hop_model,
    projection_size=1024,
    save_or_load_index=True,
    indexing_batch_size=10240,
    index_device="cuda"     # 或在 Apple Metal 上使用 cpu
)

⚠️ 重要提示

• 段落 jsonl 文件包含檢索數據庫中每個段落的 ID、標題和文本。
• 段落嵌入和 faiss 索引文件分別包含 numpy 數組格式的段落嵌入和 faiss 索引。如需替換為自己的數據庫，請參考 init_multihop_rag.py 中的邏輯。
• 有關段落文件的更詳細結構，請參考 MultiHop RAG 評估段落文件。

retriever 具有 multihop_search_passages 方法，支持單查詢和批量查詢檢索。

單查詢示例

retrieve_result = retriever.multihop_search_passages(
    "Did Engadget report a discount on the 13.6-inch MacBook Air \
        before The Verge reported a discount on Samsung Galaxy Buds 2?",
    n_hop=2,
    top_n=5
)

批量查詢示例

LIST_OF_QUESTIONS = [
    "Did Engadget report a discount on the 13.6-inch MacBook Air \
        before The Verge reported a discount on Samsung Galaxy Buds 2?",
    "Did 'The Independent - Travel' report on Tremblant Ski Resort \
        before 'Essentially Sports' mentioned Jeff Shiffrin's skiing habits?"
]

retrieve_result = retriever.multihop_search_passages(
    LIST_OF_QUESTIONS,
    n_hop=2,
    top_n=5,
    # 根據設備情況更改批量大小以優化性能
    index_batch_size=2048,
    generate_batch_size=1024
)

訪問檢索結果

# 打印檢索到的段落
print(retrieve_result.passage)

# 使用 networkx 圖描繪多跳檢索路徑
retrieve_result = retriever.multihop_search_passages(
    LIST_OF_QUESTIONS,
    n_hop=2,
    top_n=5,
    index_batch_size=2048,
    generate_batch_size=1024,
    return_tree=True        # 只需添加此參數
)
# `retrieve_result.tree_hop_graph` 是一個 networkx 對象列表
# 對應 LIST_OF_QUESTIONS 中查詢的檢索路徑。
# 以第一個查詢為例，繪製相應的路徑：
retrieval_tree = retrieve_result.tree_hop_graph[0]
retrieval_tree.plot_tree()

# 節點表示檢索圖中的段落
# 存儲原始段落的元數據：
print(retrieval_tree.nodes(data=True))

✨ 主要特性

• 處理複雜查詢：現實世界中的問題通常需要多跳檢索才能獲取相關信息，傳統檢索方法在這方面存在困難，而 TreeHop 能夠有效應對。
• 成本效益高：TreeHop 僅包含 2500 萬個參數，相比現有查詢重寫器的數十億參數，顯著降低了計算開銷。
• 速度快：與基於迭代 LLM 的方法相比，推理速度快 99%，非常適合對響應速度要求較高的工業應用。
• 性能出色：在控制檢索段落數量的情況下，保持高召回率，確保相關性的同時不會使系統負擔過重。

📦 安裝指南

本項目依賴 Python 3.9 及以上版本，具體依賴請參考 requirements.txt。

📚 詳細文檔

論文復現

要評估 TreeHop 的多跳檢索性能，請運行以下代碼。以下以 2WikiMultihop 數據集和三跳下的 recall@5 為例。腳本將打印每一跳的召回率和平均檢索段落數，以及按問題類型的統計信息。

⚠️ 重要提示

• 如需更改評估數據集，請將 2wiki 替換為 musique 或 multihop_rag。
• 修改 n_hop 和 top_n 以更改跳數和頂級檢索設置。
• 切換 redundant_pruning 和 layerwise_top_pruning 以復現我們關於停止準則的消融研究。

python evaluation.py \
    --dataset_name multihop_rag \
    --revision paper-reproduction \
    --n_hop 3 \
    --top_n 5 \
    --redundant_pruning True \
    --layerwise_top_pruning True

訓練 TreeHop

運行以下代碼生成圖並訓練 TreeHop。有關此腳本的參數，請參考 training.py 中的 parse_args 函數。有關訓練嵌入生成，請參考 init_train_vectors.py 中的代碼。

python training.py --graph_cache_dir ./train_data/

🔧 技術細節

嵌入數據庫

我們採用 BGE-m3 進行嵌入生成，並在此基礎上訓練 TreeHop 模型進行多跳檢索。運行以下兩個腳本可生成所有必要的訓練和評估嵌入數據庫。如果不打算訓練 TreeHop，則無需運行這些腳本，因為倉庫中已提供所有必要的評估嵌入數據庫。

python init_train_vectors.py
python init_multihop_rag.py

📄 許可證

本項目採用 MIT 許可證。有關詳細信息，請參閱 LICENSE。

📖 引用

如果您在研究中使用了本項目，請引用以下論文：

@misc{li2025treehopgeneratefilterquery,
      title={TreeHop: Generate and Filter Next Query Embeddings Efficiently for Multi-hop Question Answering}, 
      author={Zhonghao Li and Kunpeng Zhang and Jinghuai Ou and Shuliang Liu and Xuming Hu},
      year={2025},
      eprint={2504.20114},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2504.20114}, 
}