demo_tree_node_retrieval.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  1. # -*- coding: utf-8 -*-
  2. """
  3. 使用 BidiRag 召回 HTML 树节点内容片段示例
  4. 参考 test_bm25.py 的 test_bm25_html_tree() 方法
  5. 展示如何从 HTML 树结构中召回相关节点片段
  6. """
  7. import sys
  8. import os
  9. sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  10. from bdirag.bidi_rag import BidiRag
  11. def demo_tree_node_retrieval():
  12. """演示如何召回 HTML 树节点片段"""
  13. print("=" * 80)
  14. print("BidiRag - HTML 树节点片段召回示例")
  15. print("=" * 80)
  16. # 1. 初始化 BidiRag(使用 bm25_html_tree 方法)
  17. print("\n[步骤 1] 初始化 BidiRag (使用 bm25_html_tree 方法)...")
  18. rag = BidiRag(rag_method='bm25_html_tree')
  19. # 2. 准备 HTML 文档
  20. print("\n[步骤 2] 添加 HTML 公告...")
  21. sample_html = """
  22. <html>
  23. <body>
  24. <h1>XX市第一人民医院医疗设备招标公告</h1>
  25. <div>
  26. <h2>一、项目概况</h2>
  27. <p>项目名称:XX市第一人民医院彩色多普勒超声诊断仪采购项目</p>
  28. <p>项目编号:XX-ZB-2024-001</p>
  29. <p>预算金额:500万元</p>
  30. <p>采购内容:彩色多普勒超声诊断仪 1台</p>
  31. </div>
  32. <div>
  33. <h2>二、投标人资格要求</h2>
  34. <p>1. 具有独立承担民事责任的能力</p>
  35. <p>2. 具有有效的医疗器械经营许可证</p>
  36. <p>3. 近三年内无不良经营记录</p>
  37. <p>4. 投标保证金:人民币5万元整</p>
  38. </div>
  39. <div>
  40. <h2>三、技术需求</h2>
  41. <p>1. 彩色多普勒超声诊断仪技术参数</p>
  42. <p> - 探头配置:腹部凸阵探头、高频线阵探头、心脏相控阵探头</p>
  43. <p> - 显示屏:≥19英寸高清液晶显示器</p>
  44. <p> - 质保期:整机质保三年</p>
  45. <p>2. 交货时间:合同签订后60天内交货</p>
  46. <p>3. 交货地点:XX市第一人民医院设备科</p>
  47. </div>
  48. <div>
  49. <h2>四、评标方法</h2>
  50. <p>采用综合评分法:</p>
  51. <p> - 技术部分:60分</p>
  52. <p> - 商务部分:30分</p>
  53. <p> - 价格部分:10分</p>
  54. </div>
  55. <div>
  56. <h2>五、付款方式</h2>
  57. <p>合同签订后支付30%,交货验收合格后支付65%,质保期满后支付5%</p>
  58. </div>
  59. <div>
  60. <h2>六、投标截止时间</h2>
  61. <p>投标截止时间:2024年12月31日上午9:30</p>
  62. <p>开标时间:同投标截止时间</p>
  63. <p>投标文件递交地点:XX市公共资源交易中心</p>
  64. </div>
  65. </body>
  66. </html>
  67. """
  68. rag.add_texts([sample_html])
  69. print("✓ HTML 文档已添加")
  70. # 3. 测试不同查询,召回树节点片段
  71. test_queries = [
  72. ("预算金额", ["预算", "Budget"]),
  73. ("投标保证金", ["保证金", "Bond"]),
  74. ("技术参数 探头", ["探头", "technical"]),
  75. ("评标方法 综合评分", ["评标", "综合评分"]),
  76. ("质保期", ["质保", "Warranty"]),
  77. ("付款方式", ["付款", "Payment"]),
  78. ("交货时间", ["交货", "Delivery"]),
  79. ]
  80. print("\n" + "=" * 80)
  81. print("开始召回测试")
  82. print("=" * 80)
  83. for query, keywords in test_queries:
  84. print(f"\n{'=' * 80}")
  85. print(f"查询: {query}")
  86. print(f"{'-' * 80}")
  87. # 使用 retrieve 方法召回树节点片段
  88. results = rag.retrieve(
  89. query=query,
  90. top_k=3,
  91. keywords=keywords
  92. )
  93. print(f"召回 {len(results)} 个树节点片段:\n")
  94. # 评估相关性
  95. relevant_count = 0
  96. for i, (doc, score) in enumerate(results, 1):
  97. # 从 metadata 中获取路径信息
  98. path = doc.metadata.get("path", "")
  99. title = doc.metadata.get("title", "")
  100. # 检查相关性
  101. is_relevant = any(kw.lower() in doc.page_content.lower() for kw in keywords)
  102. if is_relevant:
  103. relevant_count += 1
  104. marker = "[✓]"
  105. else:
  106. marker = "[ ]"
  107. print(f" 片段 {i} {marker} (分数: {score:.4f})")
  108. print(f" 路径: {path}")
  109. print(f" 内容:")
  110. # 显示内容片段(前200字符)
  111. content_preview = doc.page_content[:200].replace("\n", " ")
  112. print(f" {content_preview}...")
  113. print()
  114. # 计算精确度
  115. precision = relevant_count / len(results) if results else 0
  116. print(f" 精确度@3: {precision:.1%}")
  117. print("\n" + "=" * 80)
  118. print("示例完成!")
  119. print("=" * 80)
  120. print("\n说明:")
  121. print("1. 使用 bm25_html_tree 方法可以召回 HTML 树结构的节点片段")
  122. print("2. 每个片段包含完整的路径信息(parent path)")
  123. print("3. 内容片段保留了层级结构信息")
  124. print("4. 通过 keywords 参数可以过滤和精化召回结果")
  125. if __name__ == "__main__":
  126. demo_tree_node_retrieval()