BidiRag/examples/debug_bm25.py

# -*- coding: utf-8 -*-
"""调试 BM25RAG 的 retrieve 方法"""

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

from bdirag.document_processor import Document
from bdirag.rag_methods.bm25_rag import BM25RAG

# 创建测试文档
doc = Document(
    page_content="""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市公共资源交易中心""",
    metadata={"title": "Test"}
)

# 初始化 BM25RAG
print("初始化 BM25RAG...")
bm25_rag = BM25RAG()
bm25_rag.index_documents([doc])
print(f"索引构建完成，文档数: {len(bm25_rag._all_documents)}")

# 测试查询
test_queries = [
    "预算金额",
    "投标保证金",
    "质保期",
    "评标方法",
]

for query in test_queries:
    print(f"\n{'='*60}")
    print(f"查询: {query}")
    print(f"{'='*60}")
    
    # 查看 tokenize 结果
    from bdirag.rag_methods.tokenization import bm25_tokenize
    query_tokens = bm25_tokenize(query)
    print(f"查询分词: {query_tokens}")
    
    # 检查 BM25 得分
    if bm25_rag.bm25 is not None:
        scores = bm25_rag.bm25.get_scores(query_tokens)
        print(f"BM25 原始得分: {scores}")
        print(f"最大得分: {max(scores) if len(scores) > 0 else 0:.4f}")
    
    # 执行检索
    results = bm25_rag.retrieve(query, k=3)
    print(f"召回 {len(results)} 个结果")
    
    for i, (doc, score) in enumerate(results, 1):
        print(f"  [{i}] Score: {score:.4f}")
        print(f"      Content: {doc.page_content[:100]}...")