Python之TREC：pytrec_eval简介、安装和使用方法、案例应用之详细攻略

目录

pytrec_eval的简介

pytrec_eval的安装和使用方法

1、pytrec_eval的安装

2、pytrec_eval的使用方法

pytrec_eval的案例应用

# 1、加载查询与结果，计算标准评估指标—— NDCG、MAP

# 2、批量评估多个系统

# 3、计算自定义指标

# 4、批量输出结果到文件

---

pytrec_eval的简介

pytrec_eval是一个基于流行的trec_eval的Python信息检索评估工具。它旨在避免Python编程语言中自定义信息检索评估度量实现的重复工作。它提供了一个Python接口来使用TREC的评估工具trec_eval。pytrec_eval 提供了一个方便易用的 Python 接口来进行信息检索评估，支持多种评估指标，并具有良好的扩展性，适合各种信息检索相关的研究和应用场景。

GitHub地址：GitHub - cvangysel/pytrec_eval: pytrec_eval is an Information Retrieval evaluation tool for Python, based on the popular trec_eval.

pytrec_eval的安装和使用方法

1、pytrec_eval的安装

需要一个包含开发头文件的Python 3.5或更高版本发行版，并且需要安装numpy和scipy库。安装命令如下：

pip install pytrec_eval
pip install -i https://mirrors.aliyun.com/pypi/simple pytrec_eval

2、pytrec_eval的使用方法

pytrec_eval 的核心在于 RelevanceEvaluator 类。 你需要提供查询的相关性判断 (qrel) 和运行结果 (run)。 qrel 是一个字典，键是查询ID，值是另一个字典，表示每个文档的相关性等级 (通常 0 表示不相关，1 表示相关)。 run 也是一个字典，键是查询ID，值是另一个字典，表示每个文档的得分。

import pytrec_eval
import json

qrel = {
    'q1': {
        'd1': 0,
        'd2': 1,
        'd3': 0,
    },
    'q2': {
        'd2': 1,
        'd3': 1,
    },
}

run = {
    'q1': {
        'd1': 1.0,
        'd2': 0.0,
        'd3': 1.5,
    },
    'q2': {
        'd1': 1.5,
        'd2': 0.2,
        'd3': 0.5,
    }
}

evaluator = pytrec_eval.RelevanceEvaluator(
    qrel, {'map', 'ndcg'})  # 指定需要计算的指标，这里为map和ndcg

print(json.dumps(evaluator.evaluate(run), indent=1))

这段代码会返回一个包含查询 q1 和 q2 的请求评估度量的JSON数据结构，例如：

{
    "q1": {
        "ndcg": 0.5,
        "map": 0.3333333333333333
    },
    "q2": {
        "ndcg": 0.6934264036172708,
        "map": 0.5833333333333333
    }
}

pytrec_eval的案例应用

# 1、加载查询与结果，计算标准评估指标—— NDCG、MAP

'''
pip install -i https://mirrors.aliyun.com/pypi/simple pytrec_eval
qrel 是标准答案，表示每个文档对查询的相关性得分（1 表示相关，0 表示不相关）。
run 是系统的预测得分。
evaluate 方法计算每个查询的评估指标值，例如 NDCG 和 MAP。
'''

import pytrec_eval

# 定义查询的相关性（qrel）
qrel = {
    'query1': {'doc1': 1, 'doc2': 0, 'doc3': 1},
    'query2': {'doc1': 0, 'doc2': 1},
}

# 定义系统返回的检索结果（run）
run = {
    'query1': {'doc1': 0.8, 'doc2': 0.6, 'doc3': 0.4},
    'query2': {'doc1': 0.2, 'doc2': 0.9},
}

# 初始化评估器
evaluator = pytrec_eval.RelevanceEvaluator(qrel, {'ndcg', 'map'})

# 进行评估
results = evaluator.evaluate(run)

# 输出评估结果
for query_id, query_measures in results.items():
    print(f"Query: {query_id}")
    for measure, value in query_measures.items():
        print(f"  {measure}: {value}")

# 2、批量评估多个系统

'''
假设您有多个检索系统的结果，您可以对每个系统进行批量评估。
使用字典 systems 存储多个系统的预测结果。
遍历每个系统，分别评估 NDCG 和 MAP 等指标。
'''
import pytrec_eval

# 定义查询相关性
qrel = {
    'query1': {'doc1': 1, 'doc2': 0, 'doc3': 1},
    'query2': {'doc1': 0, 'doc2': 1},
}

# 定义多个系统的结果
systems = {
    'system1': {
        'query1': {'doc1': 0.8, 'doc2': 0.6, 'doc3': 0.4},
        'query2': {'doc1': 0.2, 'doc2': 0.9},
    },
    'system2': {
        'query1': {'doc1': 0.9, 'doc2': 0.3, 'doc3': 0.7},
        'query2': {'doc1': 0.4, 'doc2': 0.8},
    },
}

# 初始化评估器
evaluator = pytrec_eval.RelevanceEvaluator(qrel, {'ndcg', 'map'})

# 对每个系统进行评估
for system_name, run in systems.items():
    results = evaluator.evaluate(run)
    print(f"Results for {system_name}:")
    for query_id, query_measures in results.items():
        print(f"  Query: {query_id}")
        for measure, value in query_measures.items():
            print(f"    {measure}: {value}")

# 3、计算自定义指标

'''
可以自定义评价指标的集合，例如只评估 precision@k。
precision@1 表示前 1 个结果的准确率。
precision@2 表示前 2 个结果的准确率。
通过指定自定义指标集合，评估更加灵活。
'''

import pytrec_eval

# 定义查询相关性
qrel = {
    'query1': {'doc1': 1, 'doc2': 0, 'doc3': 1},
    'query2': {'doc1': 0, 'doc2': 1},
}

# 定义系统返回结果
run = {
    'query1': {'doc1': 0.8, 'doc2': 0.6, 'doc3': 0.4},
    'query2': {'doc1': 0.2, 'doc2': 0.9},
}

# 初始化评估器，指定评估指标 precision@1, precision@2
evaluator = pytrec_eval.RelevanceEvaluator(qrel, {'precision@1', 'precision@2'})

# 进行评估
results = evaluator.evaluate(run)

# 输出评估结果
for query_id, query_measures in results.items():
    print(f"Query: {query_id}")
    for measure, value in query_measures.items():
        print(f"  {measure}: {value}")

# 4、批量输出结果到文件

'''
如果需要将评估结果保存到文件，可以使用以下代码：
使用 json.dump 方法将评估结果保存到文件中，方便后续分析。
'''
import pytrec_eval
import json

# 定义 qrel 和 run
qrel = {
    'query1': {'doc1': 1, 'doc2': 0, 'doc3': 1},
    'query2': {'doc1': 0, 'doc2': 1},
}

run = {
    'query1': {'doc1': 0.8, 'doc2': 0.6, 'doc3': 0.4},
    'query2': {'doc1': 0.2, 'doc2': 0.9},
}

evaluator = pytrec_eval.RelevanceEvaluator(qrel, {'ndcg', 'map'})

# 评估并保存结果
results = evaluator.evaluate(run)

# 保存到 JSON 文件
with open('results.json', 'w') as f:
    json.dump(results, f, indent=4)

print("评估结果已保存到 results.json")