test/list_benchmark_problems.py

#!/usr/bin/env python3
"""
벤치마크 문제 ID 목록 조회 도구

지원하는 벤치마크의 모든 문제 ID를 확인할 수 있습니다.
HumanEval+, MBPP+, LiveCodeBench 지원
"""

import os
import sys
import json
import argparse
from pathlib import Path
from datetime import datetime

# TestTime RLVR 모듈 임포트
sys.path.append('/home/ubuntu/RLVR/TestTime-RLVR-v2')


def load_jsonl(file_path):
    """JSONL 파일 로드"""
    if not os.path.exists(file_path):
        return []
    
    with open(file_path, 'r', encoding='utf-8') as f:
        return [json.loads(line.strip()) for line in f if line.strip()]


def list_humaneval_problems(data_path):
    """HumanEval+ 문제 목록"""
    print("🔍 HumanEval+ 문제 목록")
    print("="*60)
    
    problems = load_jsonl(data_path)
    if not problems:
        print("❌ 데이터 파일을 찾을 수 없습니다.")
        return []
    
    task_ids = []
    print(f"📊 총 {len(problems)}개 문제 발견")
    
    for i, problem in enumerate(problems):
        task_id = problem.get('task_id', f'Unknown_{i}')
        task_ids.append(task_id)
    
    # 모든 문제 ID들을 10개씩 묶어서 출력
    print("\n📋 전체 문제 ID 목록:")
    for j in range(0, len(task_ids), 10):
        batch = task_ids[j:j+10]
        print(f"  {', '.join(batch)}")
    
    return task_ids


def list_mbpp_problems(data_path):
    """MBPP+ 문제 목록"""
    print("🔍 MBPP+ 문제 목록")
    print("="*60)
    
    problems = load_jsonl(data_path)
    if not problems:
        print("❌ 데이터 파일을 찾을 수 없습니다.")
        return []
    
    task_ids = []
    print(f"📊 총 {len(problems)}개 문제 발견")
    
    for i, problem in enumerate(problems):
        task_id = problem.get('task_id', f'Unknown_{i}')
        task_ids.append(task_id)
    
    # 모든 문제 ID들을 10개씩 묶어서 출력
    print("\n📋 전체 문제 ID 목록:")
    for j in range(0, len(task_ids), 10):
        batch = task_ids[j:j+10]
        print(f"  {', '.join(batch)}")
    
    return task_ids


def list_lcb_problems(data_path):
    """LiveCodeBench 문제 목록"""
    print("🔍 LiveCodeBench 문제 목록")
    print("="*60)
    
    # LiveCodeBench는 디렉토리 구조가 다를 수 있음
    lcb_files = list(Path(data_path).glob("**/*.jsonl")) if os.path.exists(data_path) else []
    
    if not lcb_files:
        print("❌ LiveCodeBench 데이터를 찾을 수 없습니다.")
        return []
    
    all_task_ids = []
    print(f"📊 {len(lcb_files)}개 파일 발견")
    
    for file_path in lcb_files[:5]:  # 처음 5개 파일만 확인
        print(f"\n📁 파일: {file_path.name}")
        problems = load_jsonl(file_path)
        
        for i, problem in enumerate(problems[:10]):  # 각 파일에서 10개만 표시
            task_id = problem.get('task_id', problem.get('id', f'LCB_{i}'))
            all_task_ids.append(task_id)
            prompt_preview = problem.get('prompt', problem.get('description', ''))[:80].replace('\n', ' ')
            print(f"  {len(all_task_ids):3d}. {task_id} - {prompt_preview}...")
        
        if len(problems) > 10:
            print(f"      ... ({len(problems)-10}개 문제 더 있음)")
    
    if len(lcb_files) > 5:
        print(f"\n... ({len(lcb_files)-5}개 파일 더 있음)")
    
    return all_task_ids


def save_problem_list(task_ids, benchmark, output_dir):
    """문제 목록을 벤치마크별 디렉토리에 저장"""
    # tmp/{benchmark}/ 디렉토리 생성
    benchmark_dir = os.path.join(output_dir, benchmark)
    os.makedirs(benchmark_dir, exist_ok=True)
    
    # 전체 문제 목록 저장
    all_problems_file = os.path.join(benchmark_dir, f"{benchmark}_all_problems.json")
    output_data = {
        'benchmark': benchmark,
        'total_problems': len(task_ids),
        'task_ids': task_ids,
        'generated_at': datetime.now().isoformat(),
        'data_source': f'{benchmark}_plus_dataset'
    }
    
    with open(all_problems_file, 'w', encoding='utf-8') as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)
    
    print(f"\n💾 문제 목록이 저장되었습니다:")
    print(f"   전체 목록: {all_problems_file}")


def main():
    parser = argparse.ArgumentParser(description='벤치마크 문제 ID 목록 조회')
    parser.add_argument('--benchmark', type=str, default='all',
                       choices=['all', 'humaneval', 'mbpp', 'lcb'],
                       help='조회할 벤치마크 (all=모든 벤치마크)')
    parser.add_argument('--save', action='store_true',
                       help='결과를 JSON 파일로 저장')
    parser.add_argument('--output_dir', type=str, 
                       default='/home/ubuntu/RLVR/TestTime-RLVR-v2/tmp',
                       help='출력 디렉토리')
    
    args = parser.parse_args()
    
    # 데이터 경로 설정 (상대 경로로 수정)
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    humaneval_path = f'{base_dir}/evaluation/code_eval/data/HumanEvalPlus.jsonl'
    mbpp_path = f'{base_dir}/evaluation/code_eval/data/MbppPlus.jsonl'
    lcb_path = f'{base_dir}/evaluation/code_eval/coding/LiveCodeBench'
    
    os.makedirs(args.output_dir, exist_ok=True)
    
    print("🚀 TestTime RLVR 벤치마크 문제 목록 조회 도구")
    print("="*80)
    
    all_results = {}
    
    if args.benchmark in ['all', 'humaneval']:
        print("\n")
        task_ids = list_humaneval_problems(humaneval_path)
        all_results['humaneval'] = task_ids
        
        if args.save and task_ids:
            save_problem_list(task_ids, 'humaneval', args.output_dir)
    
    if args.benchmark in ['all', 'mbpp']:
        print("\n")
        task_ids = list_mbpp_problems(mbpp_path)
        all_results['mbpp'] = task_ids
        
        if args.save and task_ids:
            save_problem_list(task_ids, 'mbpp', args.output_dir)
    
    if args.benchmark in ['all', 'lcb']:
        print("\n")
        task_ids = list_lcb_problems(lcb_path)
        all_results['lcb'] = task_ids
        
        if args.save and task_ids:
            save_problem_list(task_ids, 'livecodebrench', args.output_dir)
    
    # 요약 정보
    print("\n" + "="*80)
    print("📊 벤치마크 요약")
    print("="*80)
    
    total_problems = 0
    for benchmark, task_ids in all_results.items():
        if task_ids:
            print(f"📋 {benchmark.upper()}: {len(task_ids)}개 문제")
            total_problems += len(task_ids)
            
            # 사용 가능한 문제 ID 샘플 표시
            if task_ids:
                print(f"    샘플 ID: {', '.join(task_ids[:5])}")
                if len(task_ids) > 5:
                    print(f"    ... (총 {len(task_ids)}개)")
    
    print(f"\n🎯 전체 사용 가능한 문제: {total_problems}개")
    
    # 사용법 안내
    print("\n" + "="*80)
    print("💡 사용법")
    print("="*80)
    print("테스트 실행 예시:")
    
    if 'humaneval' in all_results and all_results['humaneval']:
        sample_id = all_results['humaneval'][0]
        print(f"  python test_complete_pipeline.py --benchmark humaneval --problem_id \"{sample_id}\"")
    
    if 'mbpp' in all_results and all_results['mbpp']:
        sample_id = all_results['mbpp'][0]
        print(f"  python test_complete_pipeline.py --benchmark mbpp --problem_id \"{sample_id}\"")
    
    print("\n특정 문제만 확인:")
    print("  python list_benchmark_problems.py --benchmark mbpp")
    print("  python list_benchmark_problems.py --benchmark humaneval --save")


if __name__ == '__main__':
    main()