test_raptor_rag.py 567 B

12345678910111213141516171819
  1. # -*- coding: utf-8 -*-
  2. """RAPTORRAG standalone test."""
  3. import os
  4. import sys
  5. sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  6. from bdirag.rag_methods.raptor_rag import RAPTORRAG
  7. from examples.rag_test_utils import build_vector_rag, install_sklearn_fallback, run_retrieval_test
  8. install_sklearn_fallback()
  9. def build_tree(rag, documents):
  10. rag.build_tree(documents)
  11. if __name__ == "__main__":
  12. rag = build_vector_rag(RAPTORRAG, max_tree_depth=2, cluster_size=2)
  13. run_retrieval_test("RAPTORRAG", rag, index_func=build_tree)