run_vllm_baseline.py

#!/usr/bin/env python3
"""
使用vLLM API运行Self-Consistency Baseline评估
"""

import os
os.environ["CUDA_VISIBLE_DEVICES"]='2'
import sys
import json
import logging
import time
import argparse
from pathlib import Path
from datetime import datetime
from typing import Optional

# 添加src目录到路径
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))

from src.local_datasets import get_local_dataset_processor, sample_local_dataset
from src.vllm_models import VLLMModelWrapper, VLLMConfig, create_vllm_models_config, test_vllm_connection
from src.vllm_evaluator import VLLMSelfConsistencyEvaluator, save_vllm_results

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

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="使用vLLM API运行Self-Consistency Baseline评估")
    
    # 基本参数
    parser.add_argument("--models", nargs="+", 
                        # default=["ministral-7b", "/data/gemma-2-2b", "qwen2.5-3b"],
                        #default=["gemma-2-2b"],
                        default=["Meta-Llama-3.1-8B-Instruct"],
                        # default=["/data/Mistral-7B-Instruct-v0.3"],
                       help="要评估的模型名称列表（vLLM服务器上的模型名称）")
    parser.add_argument("--datasets", nargs="+", default=["GSM8K"], # default=["MMLU", "MMLU-Pro", "GSM8K"],
                       help="要评估的数据集列表")
    parser.add_argument("--num_samples", type=int, default=100,
                       help="每个数据集的样本数量（None表示使用全部样本）")
    

    # 新增: CoT参数
    parser.add_argument("--use_cot", default=True,
                       help="启用Chain-of-Thought prompting和答案提取")
    # 置信度分数参数 ### New
    parser.add_argument("--confidence_temp", type=float, default=0.5,
                       help="用于置信度分数softmax的温度。")
    
    # --- 新增参数 ---
    # parser.add_argument("--confidence_threshold", type=float, default=0.3,
    #                    help="置信度分数阈值。低于此值的回答将被排除在投票之外。默认为0.0（不过滤）。")
    # --- 新增参数结束 ---
    parser.add_argument("--num_confidence_prompts", type=int, default=8,
                       help="用于平均答案置信度的不同prompt的数量 (1-8)。")

    ## --- NEW/MODIFIED ARGUMENTS --- ##
    parser.add_argument("--use_path_confidence",default=True,
                       help="启用路径置信度评分。最终置信度=答案置信度*路径置信度。")
    
    parser.add_argument("--combined_confidence_threshold", type=float, default=0.3,
                       help="组合置信度分数阈值。低于此值的回答将被排除在投票之外。默认为0.0（不过滤）。")

    # Few-shot参数
    parser.add_argument("--few_shot", type=int, default=5,
                       help="Few-shot示例数量（0表示zero-shot）")
    
    # Self-consistency参数
    parser.add_argument("--n_consistency_samples", type=int, default=5,
                       help="Self-consistency采样次数")
    parser.add_argument("--temperature", type=float, default=0.6,
                       help="采样温度")
    parser.add_argument("--top_p", type=float, default=0.9,
                       help="Top-p采样参数")
    # parser.add_argument("--max_tokens", type=int, default=512,
    parser.add_argument("--max_tokens", type=int, default=512,
                       help="最大生成token数")
    # vLLM服务器配置
    parser.add_argument("--base_urls", nargs="+", 
                    #   default=["http://localhost:8079"], # gemma-2-2b
                    #    default=["http://localhost:8888"],   # ministral-7b
                        default=["http://localhost:8895"],   # Meta-Llama-3.1-8B-Instruct
                       help="vLLM服务器地址列表")
    parser.add_argument("--api_keys", nargs="+",
                       help="API密钥列表（如果需要）")
    # 输出配置
    parser.add_argument("--output_dir", default="results",
                       help="结果输出目录")
    parser.add_argument("--quick_test", action="store_true",
                       help="快速测试模式（少量样本）")
    
    return parser.parse_args()

def create_custom_vllm_config(model_key: str, base_url: str, model_name: Optional[str] = None, api_key: Optional[str] = None) -> VLLMConfig:
    """创建自定义vLLM配置"""
    return VLLMConfig(
        base_url=base_url,
        model_name=model_name or model_key,
        api_key=api_key
    )

def run_vllm_baseline_evaluation():
    """运行vLLM baseline评估"""
    args = parse_args()
    
    logger.info(f"   - 置信度温度 (T_conf): {args.confidence_temp}") ## NEW ##
    logger.info("🚀 开始使用vLLM API运行Self-Consistency Baseline评估")
    logger.info(f"📝 评估配置:")
    logger.info(f"   - 模型: {args.models}")
    logger.info(f"   - 数据集: {args.datasets}")
    logger.info(f"   - CoT模式: {'启用' if args.use_cot else '禁用'}") # 新增日志
    logger.info(f"   - Few-shot示例数: {args.few_shot}")
    logger.info(f"   - Self-consistency采样次数: {args.n_consistency_samples}")
    logger.info(f"   - Temperature: {args.temperature}")
    logger.info(f"   - Top-p: {args.top_p}")
    logger.info(f"   - 最大tokens: {args.max_tokens}")
    logger.info(f"   - 置信度Prompt数量: {args.num_confidence_prompts}") # ## NEW ##
    #logger.info(f"   - 置信度阈值: {args.confidence_threshold}") # <-- 记录新参数
    logger.info(f"   - 路径置信度评分: {'启用' if args.use_path_confidence else '禁用'}")
    if args.use_path_confidence:
        logger.info(f"   - 组合置信度阈值: {args.combined_confidence_threshold}")
    
    # 快速测试模式
    if args.quick_test:
        logger.info("🏃 启用快速测试模式")
        args.num_samples = 50  # 快速测试只用50个样本
    
    # 创建输出目录
    output_dir = Path(args.output_dir)
    output_dir.mkdir(exist_ok=True)
    
    # 测试vLLM服务器连接
    logger.info("🔗 测试vLLM服务器连接...")
    for base_url in args.base_urls:
        if not test_vllm_connection(base_url):
            logger.warning(f"⚠️ 无法连接到vLLM服务器: {base_url}")
    
    # 准备模型配置
    if args.base_urls:
        # 使用自定义配置
        vllm_configs = {}
        for i, model_key in enumerate(args.models):
            base_url = args.base_urls[i] if i < len(args.base_urls) else args.base_urls[0]
            model_name = model_key  # 统一使用models参数作为模型名称
            api_key = args.api_keys[i] if args.api_keys and i < len(args.api_keys) else None
            
            vllm_configs[model_key] = create_custom_vllm_config(
                model_key, base_url, model_name, api_key
            )
    else:
        # 使用预定义配置
        predefined_configs = create_vllm_models_config()
        vllm_configs = {k: v for k, v in predefined_configs.items() if k in args.models}
    
    # 结果存储
    all_results = {
        "timestamp": datetime.now().isoformat(),
        "evaluation_config": {
            "models": args.models,
            "datasets": args.datasets,
            "few_shot": args.few_shot,
            "n_consistency_samples": args.n_consistency_samples,
            "temperature": args.temperature,
            "top_p": args.top_p,
            "max_tokens": args.max_tokens,
            "num_samples": args.num_samples,
            "quick_test": args.quick_test,
            "use_cot": args.use_cot,
            "confidence_temp": args.confidence_temp,
            #"confidence_threshold": args.confidence_threshold # <-- 保存新参数到结果文件
            "num_confidence_prompts": args.num_confidence_prompts, # ## NEW ##
            "use_path_confidence": args.use_path_confidence,
            "combined_confidence_threshold": args.combined_confidence_threshold
        },
        "vllm_configs": {k: {"base_url": v.base_url, "model_name": v.model_name} 
                        for k, v in vllm_configs.items()},
        "results": {}
    }
    
    # 遍历每个模型和数据集
    for model_key in args.models:
        if model_key not in vllm_configs:
            logger.warning(f"⚠️ 模型 {model_key} 没有对应的vLLM配置，跳过")
            continue
        
        logger.info(f"\n🤖 开始评估模型: {model_key}")
        
        try:
            # 创建vLLM模型包装器
            config = vllm_configs[model_key]
            vllm_wrapper = VLLMModelWrapper(config)
            
            # 创建评估器
            evaluator = VLLMSelfConsistencyEvaluator(
                vllm_wrapper=vllm_wrapper,
                n_samples=args.n_consistency_samples,
                temperature=args.temperature,
                top_p=args.top_p,
                max_tokens=args.max_tokens,
                few_shot=args.few_shot,
                confidence_temp=args.confidence_temp, ## NEW ##
                use_cot=args.use_cot,  # 新增 use_cot 参数
                #confidence_threshold=args.confidence_threshold # <-- 传递新参数
                use_path_confidence=args.use_path_confidence, # Pass the flag
                combined_confidence_threshold=args.combined_confidence_threshold, # Pass the threshold
                num_confidence_prompts=args.num_confidence_prompts # ## NEW ##

            )
            
            model_results = {}
            
            # 遍历每个数据集
            for dataset_name in args.datasets:
                logger.info(f"\n📊 评估数据集: {dataset_name}")
                
                try:
                    # 加载本地数据集
                    processor = get_local_dataset_processor(dataset_name)
                    
                    # 采样数据
                    test_data = sample_local_dataset(processor, args.num_samples)
                    logger.info(f"📈 使用 {len(test_data)} 个样本进行评估")
                    
                    # 运行评估
                    start_time = time.time()
                    results = evaluator.evaluate_dataset(processor, test_data)
                    evaluation_time = time.time() - start_time
                    
                    # 保存结果
                    model_results[dataset_name] = {
                        "accuracy": results.accuracy,
                        "accuracy_percentage": results.accuracy * 100,
                        "total_samples": results.total_samples,
                        "correct_samples": results.correct_samples,
                        "avg_tokens_per_sample": results.avg_tokens_per_sample,
                        "total_tokens": results.total_tokens,
                        "avg_prompt_tokens": results.avg_prompt_tokens,
                        "avg_completion_tokens": results.avg_completion_tokens,
                        "avg_inference_time": results.avg_inference_time,
                        "total_inference_time": results.total_inference_time,
                        "num_confidence_prompts": results.num_confidence_prompts,
                        "evaluation_time": evaluation_time,
                        "tokens_per_second": results.total_tokens / results.total_inference_time if results.total_inference_time > 0 else 0
                    }
                    
                    # 保存详细结果
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    # 将模型名称中的路径分隔符替换为安全的字符
                    safe_model_name = model_key.replace("/", "_").replace("\\", "_")
                    detail_file = output_dir / f"vllm_{safe_model_name}_{dataset_name}_{timestamp}.json"
                    save_vllm_results(results, str(detail_file))
                    
                    # 打印摘要
                    logger.info(f"✅ 结果摘要 - {dataset_name}:")
                    logger.info(f"   🎯 准确率: {results.accuracy:.4f} ({results.accuracy*100:.2f}%)")
                    logger.info(f"   📝 平均tokens/问题: {results.avg_tokens_per_sample:.2f}")
                    logger.info(f"   🔢 总tokens: {results.total_tokens:,}")
                    logger.info(f"   ⏱️ 平均推理时间: {results.avg_inference_time:.2f}s")
                    logger.info(f"   🚀 吞吐量: {results.total_tokens/results.total_inference_time:.1f} tokens/s")
                    
                except Exception as e:
                    logger.error(f"❌ 评估数据集 {dataset_name} 时出错: {e}")
                    model_results[dataset_name] = {"error": str(e)}
            
            # 保存模型结果
            all_results["results"][model_key] = {
                "model_config": {
                    "base_url": config.base_url,
                    "model_name": config.model_name
                },
                "results": model_results
            }
            
        except Exception as e:
            logger.error(f"❌ 评估模型 {model_key} 时出错: {e}")
            all_results["results"][model_key] = {"error": str(e)}
    
    # 保存完整结果
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    results_file = output_dir / f"vllm_baseline_results_{timestamp}.json"
    
    with open(results_file, 'w', encoding='utf-8') as f:
        json.dump(all_results, f, indent=2, ensure_ascii=False)
    
    logger.info(f"\n💾 完整结果已保存到: {results_file}")
    
    # 生成摘要报告
    generate_summary_report(all_results, output_dir / f"vllm_baseline_summary_{timestamp}.txt")
    
    return all_results

def generate_summary_report(results: dict, output_file: Path):
    """生成摘要报告"""
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write("vLLM Self-Consistency Baseline 评估报告\n")
        f.write("=" * 50 + "\n\n")
        
        f.write(f"评估时间: {results['timestamp']}\n")
        f.write(f"评估配置: {results['evaluation_config']}\n\n")
        
        # vLLM服务器配置
        f.write("vLLM服务器配置:\n")
        for model_key, config in results["vllm_configs"].items():
            f.write(f"  - {model_key}: {config['base_url']} ({config['model_name']})\n")
        f.write("\n")
        
        # 各模型结果
        for model_key, model_data in results["results"].items():
            if "error" in model_data:
                f.write(f"模型 {model_key}: 评估失败 - {model_data['error']}\n\n")
                continue
                
            f.write(f"模型: {model_key}\n")
            f.write(f"服务器: {model_data['model_config']['base_url']}\n")
            f.write("-" * 30 + "\n")
            
            for dataset_name, dataset_results in model_data["results"].items():
                if "error" in dataset_results:
                    f.write(f"  {dataset_name}: 评估失败 - {dataset_results['error']}\n")
                    continue
                    
                accuracy = dataset_results["accuracy"]
                avg_tokens = dataset_results["avg_tokens_per_sample"]
                total_tokens = dataset_results["total_tokens"]
                throughput = dataset_results.get("tokens_per_second", 0)
                
                f.write(f"  {dataset_name}:\n")
                f.write(f"    准确率: {accuracy:.4f} ({accuracy*100:.2f}%)\n")
                f.write(f"    平均tokens/问题: {avg_tokens:.2f}\n")
                f.write(f"    总tokens: {total_tokens:,}\n")
                f.write(f"    吞吐量: {throughput:.1f} tokens/s\n")
            
            f.write("\n")
    
    logger.info(f"📋 摘要报告已保存到: {output_file}")

if __name__ == "__main__":
    try:
        results = run_vllm_baseline_evaluation()
        logger.info("🎉 vLLM Baseline评估完成！")
    except Exception as e:
        logger.error(f"💥 评估过程中出现错误: {e}")
        sys.exit(1)