BidiRag/examples/demo_bm25_retrieval.py

# -*- coding: utf-8 -*-
"""
使用 BidiRag 的 BM25 方法召回内容片段示例

对比 bm25_html_tree 和 bm25 两种方法的效果
"""

import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from bdirag.bidi_rag import BidiRag


def demo_bm25_retrieval():
    """演示使用 BM25 方法召回内容片段"""
    
    print("=" * 80)
    print("BidiRag - BM25 纯文本召回示例")
    print("=" * 80)
    
    # 1. 初始化 BidiRag（使用 bm25 方法）
    print("\n[步骤 1] 初始化 BidiRag (使用 bm25 方法)...")
    rag = BidiRag(rag_method='bm25')
    
    # 2. 准备纯文本文档（BM25 更适合纯文本）
    print("\n[步骤 2] 添加招投标公告（纯文本格式）...")
    
    sample_docs = [
        """XX市第一人民医院医疗设备招标公告

项目名称：XX市第一人民医院彩色多普勒超声诊断仪采购项目
项目编号：XX-ZB-2024-001
预算金额：500万元
采购内容：彩色多普勒超声诊断仪 1台

投标人资格要求：
1. 具有独立承担民事责任的能力
2. 具有有效的医疗器械经营许可证
3. 近三年内无不良经营记录
4. 投标保证金：人民币5万元整

技术需求：
1. 彩色多普勒超声诊断仪技术参数
   - 探头配置：腹部凸阵探头、高频线阵探头、心脏相控阵探头
   - 显示屏：≥19英寸高清液晶显示器
   - 质保期：整机质保三年
2. 交货时间：合同签订后60天内交货
3. 交货地点：XX市第一人民医院设备科

评标方法：采用综合评分法
   - 技术部分：60分
   - 商务部分：30分
   - 价格部分：10分

付款方式：合同签订后支付30%，交货验收合格后支付65%，质保期满后支付5%

投标截止时间：2024年12月31日上午9:30
开标时间：同投标截止时间
投标文件递交地点：XX市公共资源交易中心"""
    ]
    
    rag.add_texts(sample_docs)
    print("✓ 文档已添加")
    
    # 3. 测试不同查询，召回内容片段
    test_queries = [
        ("预算金额", ["预算", "Budget"]),
        ("投标保证金", ["保证金", "Bond"]),
        ("技术参数 探头", ["探头", "technical"]),
        ("评标方法 综合评分", ["评标", "综合评分"]),
        ("质保期", ["质保", "Warranty"]),
        ("付款方式", ["付款", "Payment"]),
        ("交货时间", ["交货", "Delivery"]),
    ]
    
    print("\n" + "=" * 80)
    print("开始召回测试")
    print("=" * 80)
    
    for query, keywords in test_queries:
        print(f"\n{'=' * 80}")
        print(f"查询: {query}")
        print(f"{'-' * 80}")
        
        # 使用 retrieve 方法召回内容片段
        results = rag.retrieve(
            query=query,
            top_k=3,
            keywords=keywords
        )
        
        print(f"召回 {len(results)} 个内容片段:\n")
        
        # 评估相关性
        relevant_count = 0
        for i, (doc, score) in enumerate(results, 1):
            # 检查相关性
            is_relevant = any(kw.lower() in doc.page_content.lower() for kw in keywords)
            if is_relevant:
                relevant_count += 1
                marker = "[✓]"
            else:
                marker = "[ ]"
            
            print(f"  片段 {i} {marker} (分数: {score:.4f})")
            print(f"  内容:")
            # 显示内容片段（前200字符）
            content_preview = doc.page_content[:200].replace("\n", " ")
            print(f"    {content_preview}...")
            print()
        
        # 计算精确度
        precision = relevant_count / len(results) if results else 0
        print(f"  精确度@3: {precision:.1%}")
    
    print("\n" + "=" * 80)
    print("示例完成！")
    print("=" * 80)
    print("\n说明:")
    print("1. 使用 bm25 方法召回的是纯文本片段（非树节点）")
    print("2. 片段基于文档分块，不包含层级结构信息")
    print("3. BM25 算法擅长关键词匹配")
    print("4. 适合快速检索和简单场景")


if __name__ == "__main__":
    demo_bm25_retrieval()