test_keyword_rag.py 671 B

12345678910111213141516171819202122
  1. # -*- coding: utf-8 -*-
  2. """KeywordRAG 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 examples.rag_test_utils import (
  7. FakeLLMClient,
  8. install_rank_bm25_fallback,
  9. install_sklearn_fallback,
  10. run_retrieval_test,
  11. )
  12. install_rank_bm25_fallback()
  13. install_sklearn_fallback()
  14. from bdirag.rag_methods.keyword_rag import KeywordRAG
  15. if __name__ == "__main__":
  16. run_retrieval_test("KeywordRAG-BM25", KeywordRAG(search_method="bm25", llm_client=FakeLLMClient()))
  17. print("\n")
  18. run_retrieval_test("KeywordRAG-TFIDF", KeywordRAG(search_method="tfidf", llm_client=FakeLLMClient()))