BM25_test

from rank_bm25 import BM25Okapi
from typing import Any
import asyncio


class Reranker:
    async def cut(self, text: str):
        # 模拟异步分词（实际可替换为 jieba 异步分词）
        await asyncio.sleep(0.01)
        return list(text)  # 临时按字符拆分，实际替换为 jieba.lcut(text)

    async def rerank(self, text: str,
                     corpus: list[str],
                     documents: list[Any],
                     topK: int = 10):
        # 替换为真实的异步分词
        tokenized_query = await self.cut(text)
        # 对语料库每个文本异步分词
        search_corpus = []
        for cor in corpus:
            search_corpus.append(await self.cut(cor))
        # 初始化 BM25
        bm25 = BM25Okapi(search_corpus)
        # 返回 TopK 原始文档
        return bm25.get_top_n(tokenized_query, documents, n=topK)


# 测试代码
async def main():
    reranker = Reranker()
    # 查询文本
    query = "查询工厂组织联机情况"
    # 语料库（工具描述）
    corpus = [
        "根据组织名称查询组织编码",
        "根据组织编码查询组织设备联机信息",
        "查询设备运行状态",
        "工厂人员管理"
    ]
    # 对应的原始文档（工具完整信息）
    documents = [
        {"tool_code": "AG202508220001", "tool_name": "根据组织名称查询组织编码"},
        {"tool_code": "IF20250729000006", "tool_name": "根据组织编码查询组织设备联机信息"},
        {"tool_code": "ST202508010001", "tool_name": "查询设备运行状态"},
        {"tool_code": "HR202508020001", "tool_name": "工厂人员管理"}
    ]
    # 调用重排序方法
    top3 = await reranker.rerank(query, corpus, documents, topK=1)
    print("Top3 相关文档：")
    for doc in top3:
        print(doc)


# 运行异步主函数
asyncio.run(main())


Top3 相关文档：
{'tool_code': 'HR202508020001', 'tool_name': '工厂人员管理'}
{'tool_code': 'IF20250729000006', 'tool_name': '根据组织编码查询组织设备联机信息'}
Top3 相关文档：
{'tool_code': 'HR202508020001', 'tool_name': '工厂人员管理'}
{'tool_code': 'IF20250729000006', 'tool_name': '根据组织编码查询组织设备联机信息'}
Top3 相关文档：
{'tool_code': 'HR202508020001', 'tool_name': '工厂人员管理'}
{'tool_code': 'IF20250729000006', 'tool_name': '根据组织编码查询组织设备联机信息'}
{'tool_code': 'ST202508010001', 'tool_name': '查询设备运行状态'}
Top3 相关文档：
{'tool_code': 'HR202508020001', 'tool_name': '工厂人员管理'}
Top3 相关文档：
{'tool_code': 'HR202508020001', 'tool_name': '工厂人员管理'}