from evalscope import run_task, TaskConfig
datasets = [
    # 'live_code_bench',
    'aime26',
    # 'hmmt26',
    # 'imo_answerbench',
    # 'super_gpqa',
    # 'drop',
    # 'hellaswag',
    # 'mmlu',
    # 'openai_mrcr',
    # # 'swe_bench_pro',
    # 'mcp_atlas',           
]
batch_size_list = [1, 2, 4, 8, 16, 32, 64,128]
for batch_size in batch_size_list:
    work_dir = f'/data1/sora/benchmarks/temp/output_{batch_size}'
    task_cfg = TaskConfig(
        seed=42,
        collect_perf= True,
        work_dir=work_dir,
        use_cache=work_dir,
        no_timestamp=True,
        model='DeepSeek-V4-Flash-Int8',
        api_url='http://localhost:30000/v1',
        # api_key='',
        eval_type='openai_api',
        datasets=datasets,
        dataset_dir='/data1/sora/benchmarks/bash/datasets',
        generation_config = {
            'temperature': 0.0,
            'stream': True,  
            "max_tokens": 1024 * 32,
        },
        dataset_args={
            'live_code_bench': {
                'subset_list': ['release_v6']  
            }
        },
        eval_batch_size = batch_size,
        judge_model_args = {
            "model_id": "deepseek-v4-pro",
            "api_url": "https://api.deepseek.com/v1",
            "api_key": "sk-9ed86ef546ca47e3afa7c3b014dea268",
            "eval_type": "openai_api",
            "generation_config": {
                "temperature": 0.0,
                "max_tokens": 1024 * 10,
            },
        }
    )
    run_task(task_cfg)