step4_evaluate.py

import pandas as pd
import json
import ast
import re
from tqdm import tqdm
import traceback
from collections import defaultdict
import os
import numpy as np
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.metrics import precision_recall_fscore_support, classification_report, hamming_loss

# 修改：添加函数解析多标签格式
def parse_multi_label(label_str):
    """解析多标签字符串为标签集合"""
    if isinstance(label_str, list):
        return set(label_str)
    if not isinstance(label_str, str):
        return set()
    
    try:
        # 尝试解析为JSON数组
        labels = json.loads(label_str)
        if isinstance(labels, list):
            return set(labels)
    except (json.JSONDecodeError, TypeError):
        pass
    
    # 尝试解析为Python列表字符串
    try:
        labels = ast.literal_eval(label_str)
        if isinstance(labels, list):
            return set(labels)
    except (ValueError, SyntaxError):
        pass
    
    # 尝试解析为逗号分隔的字符串
    if ',' in label_str:
        return set([tag.strip() for tag in label_str.split(',')])
    
    # 单标签情况
    return set([label_str.strip()])

# 修改：更新响应格式检查
def check_response_format(response_str):
    # 保持不变
    ...

# 修改：更新动作提取函数
def extract_action_from_response(response_str):
    if not response_str or not isinstance(response_str, str):
        return set()
    lines = [line.strip() for line in response_str.strip().splitlines() if line.strip()]
    filtered = []
    in_think = False
    for line in lines:
        if line.startswith("<think>"):
            in_think = True
            continue
        if line.endswith("</think>"):
            in_think = False
            continue
        if not in_think:
            filtered.append(line)
    for line in reversed(filtered):
        if line.startswith("Action:"):
            action_str = line[len("Action:"):].strip()
            return parse_multi_label(action_str)
    return set()

# 修改：更新真实标签提取函数
def extract_ground_truth(row):
    # 尝试从不同位置提取真实标签
    sources = [
        row.get('reward_model', {}).get('ground_truth', {}),
        row.get('extra_info', {}),
        row
    ]
    
    for source in sources:
        if isinstance(source, str):
            try:
                source = json.loads(source)
            except (TypeError, json.JSONDecodeError):
                continue
        
        if isinstance(source, dict):
            # 尝试不同键名
            for key in ['ground_truth_intent', 'ground_truth', 'answer', 'intent']:
                if key in source:
                    label_str = source[key]
                    labels = parse_multi_label(label_str)
                    if labels:
                        return labels
    
    # 尝试从特定格式提取
    if 'extra_info' in row and isinstance(row['extra_info'], str):
        match = re.search(r'<answer>(.*?)</answer>', row['extra_info'])
        if match:
            return parse_multi_label(match.group(1).strip())
    
    return set()

# 修改：重写指标计算函数为多标签版本
def calculate_multilabel_metrics(all_true, all_pred):
    """
    计算多标签分类的指标
    """
    # 获取所有唯一标签
    all_labels = sorted(set().union(*all_true, *all_pred))
    
    # 初始化二值化器
    mlb = MultiLabelBinarizer(classes=all_labels)
    y_true_bin = mlb.fit_transform(all_true)
    y_pred_bin = mlb.transform(all_pred)
    
    # 计算指标
    precision_micro, recall_micro, f1_micro, _ = precision_recall_fscore_support(
        y_true_bin, y_pred_bin, average='micro', zero_division=0
    )
    
    precision_macro, recall_macro, f1_macro, _ = precision_recall_fscore_support(
        y_true_bin, y_pred_bin, average='macro', zero_division=0
    )
    
    precision_weighted, recall_weighted, f1_weighted, _ = precision_recall_fscore_support(
        y_true_bin, y_pred_bin, average='weighted', zero_division=0
    )
    
    # 计算子集准确率
    subset_accuracy = np.mean([
        1 if true_set == pred_set else 0 
        for true_set, pred_set in zip(all_true, all_pred)
    ])
    
    # 计算汉明损失
    hamming = hamming_loss(y_true_bin, y_pred_bin)
    
    # 计算每个标签的指标
    per_label_metrics = {}
    for i, label in enumerate(all_labels):
        tp = np.sum((y_true_bin[:, i] == 1) & (y_pred_bin[:, i] == 1))
        fp = np.sum((y_true_bin[:, i] == 0) & (y_pred_bin[:, i] == 1))
        fn = np.sum((y_true_bin[:, i] == 1) & (y_pred_bin[:, i] == 0))
        tn = np.sum((y_true_bin[:, i] == 0) & (y_pred_bin[:, i] == 0))
        
        precision = tp / (tp + fp) if (tp + fp) > 0 else 0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0
        f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
        support = np.sum(y_true_bin[:, i])
        
        per_label_metrics[label] = {
            'precision': float(precision),
            'recall': float(recall),
            'f1': float(f1),
            'support': int(support),
            'tp': int(tp),
            'fp': int(fp),
            'fn': int(fn),
            'tn': int(tn)
        }
    
    # 构建指标字典
    metrics = {
        'subset_accuracy': float(subset_accuracy),
        'hamming_loss': float(hamming),
        'hamming_accuracy': float(1 - hamming),
        'micro': {
            'precision': float(precision_micro),
            'recall': float(recall_micro),
            'f1': float(f1_micro)
        },
        'macro': {
            'precision': float(precision_macro),
            'recall': float(recall_macro),
            'f1': float(f1_macro)
        },
        'weighted': {
            'precision': float(precision_weighted),
            'recall': float(recall_weighted),
            'f1': float(f1_weighted)
        },
        'per_label': per_label_metrics,
        'classification_report': classification_report(
            y_true_bin, y_pred_bin, target_names=all_labels, zero_division=0
        )
    }
    
    return metrics

# 修改：更新评估主函数
def evaluate_accuracy(parquet_path):
    print(f"Loading Parquet file: {parquet_path}")
    df = pd.read_parquet(parquet_path)
    print(f"Loaded {len(df)} rows (samples)")
    total_samples = 0
    correct_samples = 0
    sample_results = []
    error_samples = []
    correct_samples_list = []
    format_error_samples = []
    error_count = 0
    total_responses_processed = 0
    
    # 用于多标签评估的数据
    all_true_labels = []
    all_predicted_labels = []
    
    for idx, row in tqdm(df.iterrows(), total=len(df), desc="Evaluating samples"):
        try:
            ground_truth = extract_ground_truth(row)
            if not ground_truth:  # 空集合表示无法提取
                raise ValueError("Could not extract ground truth intent")
                
            responses = row['responses']
            if isinstance(responses, str):
                responses = ast.literal_eval(responses)
            elif isinstance(responses, list):
                pass
            elif hasattr(responses, 'tolist'):
                responses = responses.tolist()
            else:
                raise ValueError(f"Unexpected responses type: {type(responses)}")
                
            total_samples += 1
            total_responses_processed += len(responses)
            sample_correct = False
            sample_predictions = []
            sample_formats = []
            
            # 选择最佳预测（第一个完全匹配的，或者第一个预测）
            best_prediction = set()
            for i, resp in enumerate(responses):
                predicted = extract_action_from_response(resp)
                has_format = check_response_format(resp)
                sample_predictions.append(predicted)
                sample_formats.append(has_format)
                
                if not has_format:
                    format_error_samples.append({
                        'sample_index': idx,
                        'data_source': row.get('data_source', None),
                        'ability': row.get('ability', None),
                        'ground_truth': list(ground_truth),
                        'response': resp,
                        'predicted': list(predicted) if predicted else None
                    })
                
                # 检查是否完全匹配
                if predicted == ground_truth:
                    sample_correct = True
                    if not best_prediction:  # 只记录第一个完全匹配的
                        best_prediction = predicted
                elif not best_prediction:  # 没有完全匹配时，记录第一个预测
                    best_prediction = predicted
            
            # 记录用于多标签指标计算的标签
            all_true_labels.append(ground_truth)
            all_predicted_labels.append(best_prediction if best_prediction else set())
            
            if sample_correct:
                correct_samples += 1
                correct_samples_list.append({
                    'sample_index': idx,
                    'data_source': row.get('data_source', None),
                    'ability': row.get('ability', None),
                    'ground_truth': list(ground_truth),
                    'predictions': [list(p) for p in sample_predictions],
                    'formats': sample_formats,
                    'num_responses': len(responses),
                    'correct_predictions': [i for i, pred in enumerate(sample_predictions) if pred == ground_truth]
                })
            else:
                error_samples.append({
                    'sample_index': idx,
                    'data_source': row.get('data_source', None),
                    'ability': row.get('ability', None),
                    'ground_truth': list(ground_truth),
                    'predictions': [list(p) for p in sample_predictions],
                    'formats': sample_formats,
                    'num_responses': len(responses),
                    'correct_predictions': [i for i, pred in enumerate(sample_predictions) if pred == ground_truth]
                })
                
            sample_results.append({
                'sample_index': idx,
                'data_source': row.get('data_source', None),
                'ability': row.get('ability', None),
                'ground_truth': list(ground_truth),
                'best_prediction': list(best_prediction) if best_prediction else [],
                'is_correct': sample_correct,
                'num_responses': len(responses),
                'num_correct_responses': sum(1 for pred in sample_predictions if pred == ground_truth),
                'num_format_errors': sum(1 for f in sample_formats if not f),
                # 添加部分匹配指标
                'precision': len(ground_truth & best_prediction) / len(best_prediction) if best_prediction else 0,
                'recall': len(ground_truth & best_prediction) / len(ground_truth) if ground_truth else 0,
                'f1': 2 * len(ground_truth & best_prediction) / (len(ground_truth) + len(best_prediction)) if ground_truth or best_prediction else 0
            })
            
        except Exception as e:
            error_count += 1
            print(f"\nError processing sample {idx}: {str(e)}")
            traceback.print_exc()
            continue
    
    # 计算多标签指标
    metrics = calculate_multilabel_metrics(all_true_labels, all_predicted_labels)
    
    # 计算样本级别准确率（子集准确率）
    metrics['subset_accuracy'] = correct_samples / total_samples if total_samples > 0 else 0
    
    results_df = pd.DataFrame(sample_results)
    return metrics['subset_accuracy'], results_df, error_samples, correct_samples_list, format_error_samples, error_count, total_samples, total_responses_processed, metrics

# 其余函数保持不变
def convert_numpy_types(obj):
    ...

def get_output_paths(parquet_path):
    """
    基于输入文件路径生成输出文件路径
    """
    import os
    # 获取文件名（不含扩展名）
    base_name = os.path.splitext(os.path.basename(parquet_path))[0]
    # 创建输出目录
    output_dir = os.path.join(os.path.dirname(parquet_path), f"eval_results_{base_name}")
    os.makedirs(output_dir, exist_ok=True)
    # 生成各种输出文件路径
    output_paths = {
        'csv': os.path.join(output_dir, "sample_level_results.csv"),
        'error_json': os.path.join(output_dir, "error_samples.json"),
        'correct_json': os.path.join(output_dir, "correct_samples.json"),
        'format_error_json': os.path.join(output_dir, "format_error_samples.json"),
        'metrics_json': os.path.join(output_dir, "recall_metrics.json")
    }
    return output_paths

if __name__ == "__main__":

    # TODO: set this to your generated parquet file
    input_parquet = "/path/to/your_eval_output_data.parquet"
    sample_accuracy, detailed_results, error_samples, correct_samples_list, format_error_samples, error_count, total_samples, total_responses, metrics = evaluate_accuracy(input_parquet)
    
    # 获取输出文件路径
    output_paths = get_output_paths(input_parquet)
    
    print(f"\n{'='*50}")
    print(f"Multi-Label Evaluation Complete")
    print(f"Total Samples Processed: {total_samples}")
    print(f"Total Responses Processed: {total_responses}")
    print(f"Samples with Processing Errors: {error_count}")
    
    if not detailed_results.empty:
        print(f"Correct Samples (Exact Match): {detailed_results['is_correct'].sum()}")
        print(f"Subset Accuracy (Exact Match): {metrics['subset_accuracy']:.4f} ({metrics['subset_accuracy']*100:.2f}%)")
        print(f"Hamming Accuracy: {metrics['hamming_accuracy']:.4f} ({metrics['hamming_accuracy']*100:.2f}%)")
        print(f"Hamming Loss: {metrics['hamming_loss']:.4f}")
        
        # 打印整体指标
        print(f"\n{'='*50}")
        print(f"OVERALL METRICS")
        print(f"{'='*50}")
        print(f"Micro Precision: {metrics['micro']['precision']:.4f}")
        print(f"Micro Recall: {metrics['micro']['recall']:.4f}")
        print(f"Micro F1-Score: {metrics['micro']['f1']:.4f}")
        print(f"\nMacro Precision: {metrics['macro']['precision']:.4f}")
        print(f"Macro Recall: {metrics['macro']['recall']:.4f}")
        print(f"Macro F1-Score: {metrics['macro']['f1']:.4f}")
        print(f"\nWeighted Precision: {metrics['weighted']['precision']:.4f}")
        print(f"Weighted Recall: {metrics['weighted']['recall']:.4f}")
        print(f"Weighted F1-Score: {metrics['weighted']['f1']:.4f}")
        
        # 打印调试信息
        print(f"\n{'='*50}")
        print(f"DEBUG INFORMATION")
        print(f"{'='*50}")
        print(f"Total Samples: {total_samples}")
        
        print(f"\nPer-Label Metrics:")
        for label, details in metrics['per_label'].items():
            print(f"  {label}:")
            print(f"    Precision: {details['precision']:.4f}")
            print(f"    Recall: {details['recall']:.4f}")
            print(f"    F1: {details['f1']:.4f}")
            print(f"    Support: {details['support']}")
            print(f"    TP: {details['tp']}, FP: {details['fp']}, FN: {details['fn']}")
        
        # 打印分类报告
        print(f"\n{'='*50}")
        print(f"CLASSIFICATION REPORT")
        print(f"{'='*50}")
        print(metrics['classification_report'])
        
        if 'data_source' in detailed_results.columns:
            print(f"\nSubset Accuracy by Data Source:")
            source_accuracy = detailed_results.groupby('data_source')['is_correct'].agg(['count', 'sum', 'mean'])
            source_accuracy.columns = ['Total_Samples', 'Correct_Samples', 'Accuracy']
            print(source_accuracy)
            
            print(f"\nAverage F1 by Data Source:")
            source_f1 = detailed_results.groupby('data_source')['f1'].mean()
            print(source_f1)
            
    # 保存结果的代码保持不变...
        # 保存结果到文件
    if not detailed_results.empty:
        detailed_results.to_csv(output_paths['csv'], index=False)
        print(f"Sample-level results saved to: {output_paths['csv']}")
    else:
        print("No results to save.")
    
    if error_samples:
        with open(output_paths['error_json'], 'w', encoding='utf-8') as f:
            json.dump(error_samples, f, ensure_ascii=False, indent=2)
        print(f"Error samples saved to: {output_paths['error_json']}")
        print(f"Total error samples: {len(error_samples)}")
    else:
        print("No error samples to save.")
    
    if correct_samples_list:
        with open(output_paths['correct_json'], 'w', encoding='utf-8') as f:
            json.dump(correct_samples_list, f, ensure_ascii=False, indent=2)
        print(f"Correct samples saved to: {output_paths['correct_json']}")
        print(f"Total correct samples: {len(correct_samples_list)}")
    else:
        print("No correct samples to save.")
    
    if format_error_samples:
        with open(output_paths['format_error_json'], 'w', encoding='utf-8') as f:
            json.dump(format_error_samples, f, ensure_ascii=False, indent=2)
        print(f"Format error samples saved to: {output_paths['format_error_json']}")
        print(f"Total format error samples: {len(format_error_samples)}")
    else:
        print("No format error samples to save.")
    
    # 保存召回率指标到JSON
    with open(output_paths['metrics_json'], 'w', encoding='utf-8') as f:
        json.dump(metrics, f, ensure_ascii=False, indent=2)
    print(f"Recall metrics saved to: {output_paths['metrics_json']}")
    
    if error_samples:
        print(f"\nFound {len(error_samples)} incorrect samples. Sample errors:")
        sample_errors = error_samples[:3]
        for i, error in enumerate(sample_errors):
            print(f"\nError Sample {i+1}:")
            print(f"Sample Index: {error['sample_index']}")
            print(f"Data Source: {error['data_source']}")
            print(f"Ability: {error['ability']}")
            print(f"Ground Truth: {error['ground_truth']}")
            print(f"All Predictions: {error['predictions']}")
            print(f"Correct Predictions: {error['correct_predictions']}")
            print(f"Number of Responses: {error['num_responses']}")
            print("-"*50)
    else:
        print("\nAll samples were predicted correctly!")