evaluation/evaluate_longlamp.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LongLaMP 任务评估脚本
计算 BLEU, ROUGE, METEOR 等指标


# 指定输出文件
python evaluation/evaluate_longlamp.py \
    --predictions_file /opt/tiger/PPD/data/long_output/topic_writing_zero_shot/predictions_topic_writing.json \
    --task topic_writing \
    --output_file /opt/tiger/PPD/data/long_output/topic_writing_zero_shot_metrics.json

python evaluation/evaluate_longlamp.py \
    --predictions_file /opt/tiger/PPD/data/long_output/product_review_zero_shot/predictions_product_review.json \
    --task product_review \
    --output_file /opt/tiger/PPD/data/long_output/product_review_zero_shot_metrics.json

python evaluation/evaluate_longlamp.py \
    --predictions_file /opt/tiger/PPD/data/long_output/abstract_generation_zero_shot/predictions_abstract_generation.json \
    --task abstract_generation \
    --output_file /opt/tiger/PPD/data/long_output/abstract_generation_zero_shot_metrics.json

"""

import json
import argparse
import logging
from pathlib import Path
from typing import List, Dict
import evaluate


# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def postprocess_text(preds: List[str], labels: List[str]) -> tuple:
    """
    后处理文本：去除首尾空格
    
    Args:
        preds: 预测文本列表
        labels: 标签文本列表
        
    Returns:
        (处理后的预测列表, 处理后的标签列表)
    """
    preds = [pred.strip() for pred in preds]
    labels = [[label.strip()] for label in labels]  # ROUGE 需要列表格式
    return preds, labels


def compute_metrics(decoded_preds: List[str], decoded_labels: List[str]) -> Dict:
    """
    计算评估指标：BLEU, ROUGE, METEOR
    
    Args:
        decoded_preds: 预测文本列表
        decoded_labels: 标签文本列表
        
    Returns:
        包含所有指标的字典
    """
    # 后处理文本
    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
    
    # 加载评估指标
    logger.info("加载评估指标...")
    bleu_metric = evaluate.load("sacrebleu")
    rouge_metric = evaluate.load('rouge')
    meteor_metric = evaluate.load('meteor')
    
    # 计算 BLEU
    logger.info("计算 BLEU 分数...")
    result_bleu = bleu_metric.compute(predictions=decoded_preds, references=decoded_labels)
    
    # 计算 ROUGE
    logger.info("计算 ROUGE 分数...")
    result_rouge = rouge_metric.compute(predictions=decoded_preds, references=decoded_labels)
    
    # 计算 METEOR
    logger.info("计算 METEOR 分数...")
    result_meteor = meteor_metric.compute(predictions=decoded_preds, references=decoded_labels)
    
    # 整理结果
    result = {
        "bleu": result_bleu["score"],
        "rouge-1": result_rouge["rouge1"],
        "rouge-2": result_rouge["rouge2"],
        "rouge-L": result_rouge["rougeL"],
        "rouge-LSum": result_rouge["rougeLsum"],
        "meteor": result_meteor['meteor']
    }
    
    return result


def load_predictions(predictions_file: str) -> List[Dict]:
    """
    加载预测结果文件
    
    Args:
        predictions_file: 预测结果 JSON 文件路径
        
    Returns:
        预测结果列表
    """
    predictions_path = Path(predictions_file)
    if not predictions_path.exists():
        raise FileNotFoundError(f"预测文件不存在: {predictions_file}")
    
    with open(predictions_path, 'r', encoding='utf-8') as f:
        predictions = json.load(f)
    
    if not isinstance(predictions, list):
        raise ValueError(f"预测文件格式错误：期望列表，得到 {type(predictions)}")
    
    logger.info(f"加载了 {len(predictions)} 个预测结果")
    return predictions


def extract_predictions_and_labels(predictions: List[Dict]) -> tuple:
    """
    从预测结果中提取预测文本和标签
    
    Args:
        predictions: 预测结果列表，每个元素包含 'generated' 和 'output' 字段
        
    Returns:
        (preds, labels) 元组
    """
    preds = []
    labels = []
    skipped_count = 0
    
    for idx, pred in enumerate(predictions):
        generated = pred.get("generated", "")
        output = pred.get("output", "")
        
        # 只处理有效的预测和标签
        if generated and output:
            preds.append(generated)
            labels.append(output)
        else:
            skipped_count += 1
            sample_id = pred.get('id', f'index_{idx}')
            logger.warning(
                f"跳过无效样本: id={sample_id}, "
                f"generated={'存在' if generated else '缺失'}, "
                f"output={'存在' if output else '缺失'}"
            )
    
    if skipped_count > 0:
        logger.warning(f"跳过了 {skipped_count} 个无效样本")
    
    logger.info(f"提取了 {len(preds)} 个有效预测-标签对")
    
    if len(preds) == 0:
        raise ValueError("没有有效的预测-标签对，无法进行评估")
    
    return preds, labels


def evaluate_predictions(
    predictions_file: str,
    output_file: str = None,
    task_name: str = None
) -> Dict:
    """
    评估预测结果
    
    Args:
        predictions_file: 预测结果 JSON 文件路径
        output_file: 输出结果文件路径（可选）
        task_name: 任务名称（可选，用于输出文件名）
        
    Returns:
        评估结果字典
    """
    logger.info("=" * 50)
    logger.info("开始评估")
    logger.info("=" * 50)
    logger.info(f"预测文件: {predictions_file}")
    
    # 加载预测结果
    predictions = load_predictions(predictions_file)
    
    # 提取预测和标签
    preds, labels = extract_predictions_and_labels(predictions)
    
    # 计算指标
    logger.info("计算评估指标...")
    metrics = compute_metrics(preds, labels)
    
    # 打印结果
    logger.info("=" * 50)
    logger.info("评估结果")
    logger.info("=" * 50)
    for metric_name, metric_value in metrics.items():
        if isinstance(metric_value, dict):
            # 某些指标返回字典（如 ROUGE）
            logger.info(f"{metric_name}:")
            for k, v in metric_value.items():
                if isinstance(v, (int, float)):
                    logger.info(f"  {k}: {v:.4f}")
                else:
                    logger.info(f"  {k}: {v}")
        else:
            if isinstance(metric_value, (int, float)):
                logger.info(f"{metric_name}: {metric_value:.4f}")
            else:
                logger.info(f"{metric_name}: {metric_value}")
    
    # 保存结果
    if output_file:
        output_path = Path(output_file)
    else:
        # 如果没有指定输出文件，使用默认路径
        predictions_path = Path(predictions_file)
        if task_name:
            output_path = predictions_path.parent / f"metrics_{task_name}.json"
        else:
            output_path = predictions_path.parent / "metrics.json"
    
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    # 构建完整结果
    result = {
        "task": task_name or "unknown",
        "num_samples": len(preds),
        "metrics": metrics
    }
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    
    logger.info(f"评估结果已保存到: {output_path}")
    
    return metrics


def main():
    parser = argparse.ArgumentParser(
        description="LongLaMP 任务评估脚本",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  # 基本用法
  python evaluate_longlamp.py \\
      --predictions_file ./output/predictions_topic_writing.json \\
      --task topic_writing
  
  # 指定输出文件
  python evaluate_longlamp.py \\
      --predictions_file ./output/predictions_topic_writing.json \\
      --task topic_writing \\
      --output_file ./output/metrics.json
        """
    )
    
    parser.add_argument(
        "--predictions_file",
        type=str,
        required=True,
        help="预测结果 JSON 文件路径（推理脚本的输出文件）"
    )
    parser.add_argument(
        "--output_file",
        type=str,
        default=None,
        help="评估结果输出文件路径（可选，默认保存在预测文件同目录）"
    )
    parser.add_argument(
        "--task",
        type=str,
        default=None,
        choices=["abstract_generation", "product_review", "topic_writing"],
        help="任务名称（可选，用于输出文件名）"
    )
    
    args = parser.parse_args()
    
    try:
        # 运行评估
        metrics = evaluate_predictions(
            predictions_file=args.predictions_file,
            output_file=args.output_file,
            task_name=args.task
        )
        
        logger.info("=" * 50)
        logger.info("评估完成")
        logger.info("=" * 50)
        
        # 返回退出码
        return 0
        
    except FileNotFoundError as e:
        logger.error(f"文件未找到: {e}")
        return 1
    except ValueError as e:
        logger.error(f"数据格式错误: {e}")
        return 1
    except Exception as e:
        logger.error(f"评估过程中出错: {e}", exc_info=True)
        return 1


if __name__ == "__main__":
    exit(main())