inference/evaluate_testing.py

import json
import numpy as np
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from rouge_score import rouge_scorer
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
import nltk
from tqdm import tqdm
import os
from typing import List, Dict, Any
import argparse
import matplotlib.pyplot as plt
import seaborn as sns
from collections import defaultdict

class AnswerEvaluator:
    def __init__(self):
        try:
            nltk.data.find('tokenizers/punkt')
        except LookupError:
            nltk.download('punkt')
        
        # 初始化ROUGE评分器
        self.rouge_scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True)
        self.smoothing = SmoothingFunction().method1
    
    def load_results(self, results_path: str) -> List[Dict[str, Any]]:
        """加载生成的结果"""
        print(f"Loading results from {results_path}")
        
        with open(results_path, 'r', encoding='utf-8') as f:
            results = json.load(f)
        
        print(f"Loaded {len(results)} results")
        return results
    
    def calculate_rouge_scores(self, generated: str, reference: str) -> Dict[str, float]:
        """计算ROUGE分数"""
        try:
            scores = self.rouge_scorer.score(reference, generated)
            return {
                'rouge1_f1': scores['rouge1'].fmeasure,
                'rouge2_f1': scores['rouge2'].fmeasure,
                'rougeL_f1': scores['rougeL'].fmeasure,
                'rouge1_precision': scores['rouge1'].precision,
                'rouge2_precision': scores['rouge2'].precision,
                'rougeL_precision': scores['rougeL'].precision,
                'rouge1_recall': scores['rouge1'].recall,
                'rouge2_recall': scores['rouge2'].recall,
                'rougeL_recall': scores['rougeL'].recall
            }
        except Exception as e:
            print(f"Error calculating ROUGE scores: {e}")
            return {
                'rouge1_f1': 0.0, 'rouge2_f1': 0.0, 'rougeL_f1': 0.0,
                'rouge1_precision': 0.0, 'rouge2_precision': 0.0, 'rougeL_precision': 0.0,
                'rouge1_recall': 0.0, 'rouge2_recall': 0.0, 'rougeL_recall': 0.0
            }
    
    def calculate_bleu_score(self, generated: str, reference: str) -> float:
        """计算BLEU分数"""
        try:
            # 分词
            generated_tokens = nltk.word_tokenize(generated.lower())
            reference_tokens = nltk.word_tokenize(reference.lower())
            
            # 计算BLEU分数
            bleu_score = sentence_bleu([reference_tokens], generated_tokens, 
                                     smoothing_function=self.smoothing)
            return bleu_score
        except Exception as e:
            print(f"Error calculating BLEU score: {e}")
            return 0.0
    
    def calculate_length_metrics(self, generated: str, reference: str) -> Dict[str, float]:
        """计算长度相关指标"""
        gen_words = len(generated.split())
        ref_words = len(reference.split())
        
        return {
            'generated_length': gen_words,
            'reference_length': ref_words,
            'length_ratio': gen_words / max(ref_words, 1),
            'length_difference': abs(gen_words - ref_words)
        }
    
    def calculate_semantic_similarity(self, generated: str, reference: str) -> float:
        """计算语义相似度（简单的词汇重叠）"""
        try:
            gen_words = set(generated.lower().split())
            ref_words = set(reference.lower().split())
            
            if not gen_words or not ref_words:
                return 0.0
            
            intersection = len(gen_words.intersection(ref_words))
            union = len(gen_words.union(ref_words))
            
            return intersection / union if union > 0 else 0.0
        except Exception as e:
            print(f"Error calculating semantic similarity: {e}")
            return 0.0
    
    def evaluate_single_answer(self, generated: str, reference: str) -> Dict[str, Any]:
        """评估单个答案"""
        # 确保字符串不为空
        generated = generated.strip() if generated else ""
        reference = reference.strip() if reference else ""
        
        # 计算各种指标
        rouge_scores = self.calculate_rouge_scores(generated, reference)
        bleu_score = self.calculate_bleu_score(generated, reference)
        length_metrics = self.calculate_length_metrics(generated, reference)
        semantic_similarity = self.calculate_semantic_similarity(generated, reference)
        
        return {
            **rouge_scores,
            'bleu_score': bleu_score,
            **length_metrics,
            'semantic_similarity': semantic_similarity,
            'generated_text': generated,
            'reference_text': reference
        }
    
    def evaluate_all_answers(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """评估所有答案"""
        print("Evaluating all answers...")
        
        evaluations = []
        cluster_evaluations = defaultdict(list)
        
        for result in tqdm(results, desc="Evaluating"):
            generated = result.get('generated_answer', '')
            reference = result.get('ground_truth', '')
            cluster_id = result.get('nearest_cluster', 0)
            
            # 评估单个答案
            evaluation = self.evaluate_single_answer(generated, reference)
            evaluation['user_id'] = result.get('user_id', '')
            evaluation['cluster_id'] = cluster_id
            evaluation['query'] = result.get('query', '')
            
            evaluations.append(evaluation)
            cluster_evaluations[cluster_id].append(evaluation)
        
        # 计算总体统计
        overall_stats = self.calculate_overall_statistics(evaluations)
        
        # 计算每个聚类的统计
        cluster_stats = {}
        for cluster_id, cluster_evals in cluster_evaluations.items():
            cluster_stats[cluster_id] = self.calculate_overall_statistics(cluster_evals)
        
        return {
            'overall_statistics': overall_stats,
            'cluster_statistics': cluster_stats,
            'individual_evaluations': evaluations
        }
    
    def calculate_overall_statistics(self, evaluations: List[Dict[str, Any]]) -> Dict[str, Any]:
        """计算总体统计"""
        if not evaluations:
            return {}
        
        # 提取数值指标
        metrics = ['rouge1_f1', 'rouge2_f1', 'rougeL_f1', 'bleu_score', 
                  'semantic_similarity', 'length_ratio']
        
        stats = {}
        for metric in metrics:
            values = [eval_[metric] for eval_ in evaluations if metric in eval_]
            if values:
                stats[f'{metric}_mean'] = np.mean(values)
                stats[f'{metric}_std'] = np.std(values)
                stats[f'{metric}_min'] = np.min(values)
                stats[f'{metric}_max'] = np.max(values)
        
        # 计算长度统计
        gen_lengths = [eval_['generated_length'] for eval_ in evaluations]
        ref_lengths = [eval_['reference_length'] for eval_ in evaluations]
        
        stats['avg_generated_length'] = np.mean(gen_lengths)
        stats['avg_reference_length'] = np.mean(ref_lengths)
        stats['total_samples'] = len(evaluations)
        
        return stats
    
    def create_visualizations(self, evaluation_results: Dict[str, Any], output_dir: str):
        """创建可视化图表"""
        print("Creating visualizations...")
        
        os.makedirs(output_dir, exist_ok=True)
        
        # 1. 总体性能雷达图
        self.create_radar_chart(evaluation_results['overall_statistics'], 
                               os.path.join(output_dir, 'overall_performance.png'))
        
        # 2. 聚类性能对比
        self.create_cluster_comparison(evaluation_results['cluster_statistics'], 
                                     os.path.join(output_dir, 'cluster_performance.png'))
        
        # 3. 指标分布图
        self.create_metric_distributions(evaluation_results['individual_evaluations'], 
                                       os.path.join(output_dir, 'metric_distributions.png'))
    
    def create_radar_chart(self, stats: Dict[str, Any], save_path: str):
        """创建雷达图"""
        metrics = ['rouge1_f1', 'rouge2_f1', 'rougeL_f1', 'bleu_score', 'semantic_similarity']
        values = [stats.get(f'{m}_mean', 0) for m in metrics]
        
        # 添加第一个值到末尾以闭合图形
        values += values[:1]
        angles = np.linspace(0, 2 * np.pi, len(metrics), endpoint=False).tolist()
        angles += angles[:1]
        
        fig, ax = plt.subplots(figsize=(8, 8), subplot_kw=dict(projection='polar'))
        ax.plot(angles, values, 'o-', linewidth=2)
        ax.fill(angles, values, alpha=0.25)
        ax.set_xticks(angles[:-1])
        ax.set_xticklabels(metrics)
        ax.set_ylim(0, 1)
        ax.set_title('Overall Performance Metrics', size=16, pad=20)
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
    
    def create_cluster_comparison(self, cluster_stats: Dict[str, Any], save_path: str):
        """创建聚类性能对比图"""
        if not cluster_stats:
            return
        
        metrics = ['rouge1_f1', 'rouge2_f1', 'rougeL_f1', 'bleu_score']
        clusters = list(cluster_stats.keys())
        
        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
        axes = axes.ravel()
        
        for i, metric in enumerate(metrics):
            values = [cluster_stats[cluster].get(f'{metric}_mean', 0) for cluster in clusters]
            
            axes[i].bar(clusters, values)
            axes[i].set_title(f'{metric.upper()} by Cluster')
            axes[i].set_ylabel('Score')
            axes[i].set_ylim(0, 1)
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
    
    def create_metric_distributions(self, evaluations: List[Dict[str, Any]], save_path: str):
        """创建指标分布图"""
        metrics = ['rouge1_f1', 'rouge2_f1', 'rougeL_f1', 'bleu_score', 'semantic_similarity']
        
        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
        axes = axes.ravel()
        
        for i, metric in enumerate(metrics):
            values = [eval_[metric] for eval_ in evaluations if metric in eval_]
            
            axes[i].hist(values, bins=20, alpha=0.7, edgecolor='black')
            axes[i].set_title(f'{metric.upper()} Distribution')
            axes[i].set_xlabel('Score')
            axes[i].set_ylabel('Frequency')
            axes[i].axvline(np.mean(values), color='red', linestyle='--', 
                           label=f'Mean: {np.mean(values):.3f}')
            axes[i].legend()
        
        # 隐藏多余的子图
        axes[-1].set_visible(False)
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
    
    def print_evaluation_summary(self, evaluation_results: Dict[str, Any]):
        """打印评估摘要"""
        overall_stats = evaluation_results['overall_statistics']
        cluster_stats = evaluation_results['cluster_statistics']
        
        print("\n" + "="*60)
        print("📊 EVALUATION SUMMARY")
        print("="*60)
        
        print(f"\n📈 Overall Performance:")
        print(f"  Total samples: {overall_stats.get('total_samples', 0)}")
        print(f"  ROUGE-1 F1: {overall_stats.get('rouge1_f1_mean', 0):.4f}")
        print(f"  ROUGE-2 F1: {overall_stats.get('rouge2_f1_mean', 0):.4f}")
        print(f"  ROUGE-L F1: {overall_stats.get('rougeL_f1_mean', 0):.4f}")
        print(f"  BLEU Score: {overall_stats.get('bleu_score_mean', 0):.4f}")
        print(f"  Semantic Similarity: {overall_stats.get('semantic_similarity_mean', 0):.4f}")
        
        print(f"\n📏 Length Statistics:")
        print(f"  Avg Generated Length: {overall_stats.get('avg_generated_length', 0):.1f}")
        print(f"  Avg Reference Length: {overall_stats.get('avg_reference_length', 0):.1f}")
        
        print(f"\n🎯 Cluster Performance:")
        for cluster_id, stats in cluster_stats.items():
            print(f"  Cluster {cluster_id}:")
            print(f"    Samples: {stats.get('total_samples', 0)}")
            print(f"    ROUGE-1 F1: {stats.get('rouge1_f1_mean', 0):.4f}")
            print(f"    BLEU Score: {stats.get('bleu_score_mean', 0):.4f}")
    
    def save_evaluation_results(self, evaluation_results: Dict[str, Any], output_path: str):
        """保存评估结果"""
        print(f"Saving evaluation results to {output_path}")
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(evaluation_results, f, indent=2, ensure_ascii=False, default=str)

def main():
    parser = argparse.ArgumentParser(description='评估生成的答案')
    parser.add_argument('--results_path', type=str, required=True,
                       help='生成结果文件路径')
    parser.add_argument('--output_dir', type=str, default='./evaluation_results',
                       help='输出目录')
    parser.add_argument('--create_visualizations', action='store_true',
                       help='是否创建可视化图表')
    
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 初始化评估器
    evaluator = AnswerEvaluator()
    
    # 加载结果
    results = evaluator.load_results(args.results_path)
    
    # 评估答案
    evaluation_results = evaluator.evaluate_all_answers(results)
    
    # 打印摘要
    evaluator.print_evaluation_summary(evaluation_results)
    
    # 保存结果
    results_path = os.path.join(args.output_dir, 'evaluation_results.json')
    evaluator.save_evaluation_results(evaluation_results, results_path)
    
    # 创建可视化
    if args.create_visualizations:
        viz_dir = os.path.join(args.output_dir, 'visualizations')
        evaluator.create_visualizations(evaluation_results, viz_dir)
    
    print(f"\n✅ Evaluation completed!")
    print(f"Results saved to: {args.output_dir}")

if __name__ == "__main__":
    main()