eval.py

"""
评测模块: 用测试集评估检索质量

三个指标:
- Hit Rate: 正确文件是否出现在检索结果中(0 或 1)
- Context Precision: 检索结果中来自正确文件的比例
- Context Recall: 期望关键词在检索结果文本中出现的比例

使用方式: python eval.py
"""

import json
import os

from document_loader import load_documents
from text_splitter import split_documents
from vector_store import get_collection, add_documents, clear_store
from retriever import BM25Index, hybrid_search


# ============ 配置 ============
DOCS_DIR = "F:/AI_Program/enterprise-kb/test_docs"
EVAL_DATASET = "F:/AI_Program/enterprise-kb/eval_dataset.json"
TOP_K = 5


# ============ 评测指标计算 ============
def calc_hit_rate(results, expected_file: str) -> float:
    """
    命中率: 检索结果中是否至少有一个 chunk 来自期望文件.

    返回 1.0(命中) 或 0.0(未命中)
    """
    for r in results:
        if r.metadata.get("filename", "") == expected_file:
            return 1.0
    return 0.0


def calc_precision(results, expected_file: str) -> float:
    """
    上下文精度: 检索结果中来自正确文件的 chunk 占比.

    例: 返回 5 条, 3 条来自正确文件 -> 3/5 = 0.6
    """
    if not results:
        return 0.0

    correct = sum(
        1 for r in results
        if r.metadata.get("filename", "") == expected_file
    )
    return correct / len(results)


def calc_recall(results, expected_keywords: list[str]) -> float:
    """
    上下文召回: 期望关键词在检索结果文本中出现的比例.

    把所有检索到的 chunk 内容拼起来, 检查每个关键词是否出现.
    例: 3 个关键词, 命中 2 个 -> 2/3 = 0.67
    """
    if not expected_keywords:
        return 1.0

    # 所有检索结果拼成一个大文本, 统一小写做匹配
    all_text = " ".join(r.content for r in results).lower()

    hit = sum(
        1 for keyword in expected_keywords
        if keyword.lower() in all_text
    )
    return hit / len(expected_keywords)


# ============ 主评测流程 ============
def run_eval(
    docs_dir: str = DOCS_DIR,
    dataset_path: str = EVAL_DATASET,
    top_k: int = TOP_K,
    use_rerank: bool = True,
) -> dict:
    """
    执行完整评测流程.

    1. 加载文档并建索引
    2. 读取测试集
    3. 对每条测试数据执行检索并打分
    4. 汇总输出报告

    返回:
        包含总体指标和每条详情的字典
    """
    # --- 第一步: 建索引(跟 main.py 的 init_index 一样) ---
    print(f"\n{'='*50}")
    print("评测模式: 正在建索引...")
    print(f"{'='*50}")

    clear_store()
    docs = load_documents(docs_dir)
    if not docs:
        print("没有文档, 无法评测")
        return {}

    chunks = split_documents(docs, chunk_size=500, chunk_overlap=100)
    collection = get_collection()
    add_documents(chunks, collection)

    bm25_index = BM25Index()
    bm25_index.build(chunks)

    # --- 第二步: 读取测试集 ---
    with open(dataset_path, "r", encoding="utf-8") as f:
        dataset = json.load(f)

    print(f"\n测试集: {len(dataset)} 条问题")
    print(f"检索参数: top_k={top_k}, use_rerank={use_rerank}")
    print(f"{'='*50}\n")

    # --- 第三步: 逐条评测 ---
    all_hit_rates = []
    all_precisions = []
    all_recalls = []
    details = []

    for i, item in enumerate(dataset):
        question = item["question"]
        expected_file = item["expected_file"]
        expected_keywords = item["expected_keywords"]

        # 执行检索
        results = hybrid_search(
            query=question,
            bm25_index=bm25_index,
            collection=collection,
            k=top_k,
            use_rerank=use_rerank,
        )

        # 计算三个指标
        hit_rate = calc_hit_rate(results, expected_file)
        precision = calc_precision(results, expected_file)
        recall = calc_recall(results, expected_keywords)

        all_hit_rates.append(hit_rate)
        all_precisions.append(precision)
        all_recalls.append(recall)

        # 记录详情
        result_files = [r.metadata.get("filename", "?") for r in results]
        detail = {
            "question": question,
            "expected_file": expected_file,
            "expected_keywords": expected_keywords,
            "result_files": result_files,
            "hit_rate": hit_rate,
            "precision": precision,
            "recall": recall,
        }
        details.append(detail)

        # 打印每条结果
        status = "HIT" if hit_rate > 0 else "MISS"
        print(f"  [{status}] Q{i+1}: {question}")
        print(f"         期望: {expected_file}")
        print(f"         实际: {result_files}")
        print(f"         精度={precision:.2f}  召回={recall:.2f}")
        print()

    # --- 第四步: 汇总 ---
    avg_hit_rate = sum(all_hit_rates) / len(all_hit_rates)
    avg_precision = sum(all_precisions) / len(all_precisions)
    avg_recall = sum(all_recalls) / len(all_recalls)

    print(f"{'='*50}")
    print(f"评测报告 (共 {len(dataset)} 条)")
    print(f"{'='*50}")
    print(f"  Hit Rate        : {avg_hit_rate:.2%}  ({sum(all_hit_rates):.0f}/{len(all_hit_rates)} 命中)")
    print(f"  Context Precision: {avg_precision:.2%}")
    print(f"  Context Recall   : {avg_recall:.2%}")
    print(f"{'='*50}")

    return {
        "hit_rate": avg_hit_rate,
        "precision": avg_precision,
        "recall": avg_recall,
        "total_questions": len(dataset),
        "details": details,
    }


# ============ 对比评测: 不同配置的效果差异 ============
def compare_configs(
    docs_dir: str = DOCS_DIR,
    dataset_path: str = EVAL_DATASET,
    top_k: int = TOP_K,
) -> None:
    """
    对比两种配置的检索效果:
    1. 不用 Rerank(只有 BM25 + 向量 + RRF)
    2. 用 Rerank(加 Cross-Encoder 精排)

    只建一次索引, 避免 ChromaDB 文件锁冲突.
    """
    print("\n" + "=" * 60)
    print("配置对比: Rerank OFF vs Rerank ON")
    print("=" * 60)

    # 只建一次索引, 两次评测复用
    print(f"\n{'='*50}")
    print("评测模式: 正在建索引...")
    print(f"{'='*50}")

    clear_store()
    docs = load_documents(docs_dir)
    if not docs:
        print("没有文档, 无法评测")
        return

    chunks = split_documents(docs, chunk_size=500, chunk_overlap=100)
    collection = get_collection()
    add_documents(chunks, collection)

    bm25_index = BM25Index()
    bm25_index.build(chunks)

    with open(dataset_path, "r", encoding="utf-8") as f:
        dataset = json.load(f)

    # 两种配置分别评测
    configs = [
        ("Rerank OFF", False),
        ("Rerank ON", True),
    ]
    all_results = {}

    for config_name, use_rerank in configs:
        print(f"\n--- {config_name} ---")
        print(f"测试集: {len(dataset)} 条问题, top_k={top_k}, use_rerank={use_rerank}\n")

        hit_rates, precisions, recalls = [], [], []

        for i, item in enumerate(dataset):
            results = hybrid_search(
                query=item["question"],
                bm25_index=bm25_index,
                collection=collection,
                k=top_k,
                use_rerank=use_rerank,
            )

            hr = calc_hit_rate(results, item["expected_file"])
            pr = calc_precision(results, item["expected_file"])
            rc = calc_recall(results, item["expected_keywords"])
            hit_rates.append(hr)
            precisions.append(pr)
            recalls.append(rc)

            status = "HIT" if hr > 0 else "MISS"
            result_files = [r.metadata.get("filename", "?") for r in results]
            print(f"  [{status}] Q{i+1}: {item['question']}")
            print(f"         实际: {result_files}  精度={pr:.2f}  召回={rc:.2f}")

        avg_hr = sum(hit_rates) / len(hit_rates)
        avg_pr = sum(precisions) / len(precisions)
        avg_rc = sum(recalls) / len(recalls)

        print(f"\n  Hit Rate: {avg_hr:.2%} | Precision: {avg_pr:.2%} | Recall: {avg_rc:.2%}")

        all_results[config_name] = {"hit_rate": avg_hr, "precision": avg_pr, "recall": avg_rc}

    # 输出对比表
    print("\n" + "=" * 60)
    print("对比结果")
    print("=" * 60)
    print(f"  {'指标':<20} {'Rerank OFF':>12} {'Rerank ON':>12} {'提升':>10}")
    print(f"  {'-'*54}")

    for metric, name in [("hit_rate", "Hit Rate"), ("precision", "Precision"), ("recall", "Recall")]:
        off = all_results["Rerank OFF"][metric]
        on = all_results["Rerank ON"][metric]
        diff = on - off
        sign = "+" if diff >= 0 else ""
        print(f"  {name:<20} {off:>11.2%} {on:>11.2%} {sign}{diff:>9.2%}")

    print(f"  {'-'*54}")


# ============ 启动 ============
if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "compare":
        # python eval.py compare → 对比模式
        compare_configs()
    else:
        # python eval.py → 默认评测(开 Rerank)
        run_eval()