run_evaluation.py

import os
import argparse
import json
from datetime import datetime
from typing import Dict, Any

from src.models import get_model, InferenceConfig
from src.datasets import get_dataset_processor, sample_dataset
from src.evaluator import SelfConsistencyEvaluator, save_results
from src.analysis import create_summary_report, plot_results

import logging

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('evaluation.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

def main():
    parser = argparse.ArgumentParser(description="Self-Consistency Baseline Evaluation")
    
    parser.add_argument("--models", nargs="+", 
                       default=["ministral-7b", "gemma-2-2b", "qwen2.5-3b"],
                       help="Models to evaluate")
    
    parser.add_argument("--datasets", nargs="+",
                       default=["mmlu-pro", "gsm8k"],
                       help="Datasets to evaluate on")
    
    parser.add_argument("--num_samples", type=int, default=None,
                       help="Number of samples to evaluate (None for all)")
    
    parser.add_argument("--output_dir", type=str, default="results",
                       help="Output directory for results")
    
    parser.add_argument("--temperature", type=float, default=0.7,
                       help="Sampling temperature")
    
    parser.add_argument("--top_p", type=float, default=0.9,
                       help="Top-p sampling parameter")
    
    parser.add_argument("--num_consistency_samples", type=int, default=5,
                       help="Number of samples for self-consistency")
    
    parser.add_argument("--max_new_tokens", type=int, default=512,
                       help="Maximum new tokens to generate")
    
    args = parser.parse_args()
    
    # 创建输出目录
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = os.path.join(args.output_dir, f"evaluation_{timestamp}")
    os.makedirs(output_dir, exist_ok=True)
    
    # 保存配置
    config = {
        "models": args.models,
        "datasets": args.datasets,
        "num_samples": args.num_samples,
        "inference_config": {
            "num_samples": args.num_consistency_samples,
            "temperature": args.temperature,
            "top_p": args.top_p,
            "max_new_tokens": args.max_new_tokens
        },
        "timestamp": timestamp
    }
    
    with open(os.path.join(output_dir, "config.json"), "w") as f:
        json.dump(config, f, indent=2)
    
    # 创建推理配置
    inference_config = InferenceConfig(
        num_samples=args.num_consistency_samples,
        temperature=args.temperature,
        top_p=args.top_p,
        max_new_tokens=args.max_new_tokens
    )
    
    all_results = {}
    
    # 评估每个模型在每个数据集上的表现
    for model_name in args.models:
        logger.info(f"Starting evaluation for model: {model_name}")
        
        try:
            # 加载模型
            model = get_model(model_name)
            evaluator = SelfConsistencyEvaluator(model, inference_config)
            
            model_results = {}
            
            for dataset_name in args.datasets:
                logger.info(f"Evaluating {model_name} on {dataset_name}")
                
                try:
                    # 加载数据集
                    processor = get_dataset_processor(dataset_name)
                    
                    # 运行评估
                    results = evaluator.evaluate_dataset(processor, args.num_samples)
                    
                    # 保存单个结果
                    result_file = os.path.join(output_dir, f"{model_name}_{dataset_name}_results.json")
                    save_results(results, result_file)
                    
                    model_results[dataset_name] = results
                    
                    logger.info(f"Completed {model_name} on {dataset_name}: "
                               f"Accuracy = {results.accuracy:.4f}, "
                               f"Avg Tokens = {results.avg_tokens_per_sample:.1f}")
                    
                except Exception as e:
                    logger.error(f"Error evaluating {model_name} on {dataset_name}: {e}")
                    continue
            
            all_results[model_name] = model_results
            
        except Exception as e:
            logger.error(f"Error loading model {model_name}: {e}")
            continue
    
    # 生成汇总报告
    if all_results:
        logger.info("Generating summary report...")
        
        try:
            # 创建汇总报告
            summary_file = os.path.join(output_dir, "summary_report.md")
            create_summary_report(all_results, summary_file)
            
            # 创建可视化图表
            plot_file = os.path.join(output_dir, "results_visualization.png")
            plot_results(all_results, plot_file)
            
            logger.info(f"Evaluation completed. Results saved in {output_dir}")
            
        except Exception as e:
            logger.error(f"Error generating summary: {e}")
    
    else:
        logger.error("No successful evaluations completed")

if __name__ == "__main__":
    main()