batch_test.py

#!/usr/bin/env python3
"""
Batch test script - Run test mode with v5~v8 segment files
自動批次測試腳本 - 使用 v5~v8 segment 檔案執行測試

Usage:
    python batch_test.py --mode test --llm-provider gemini
    python batch_test.py --mode train --llm-provider gemini
"""

import os
import sys
import subprocess
import argparse
from datetime import datetime

# Segment versions to test (version will be auto-incremented by config.py)
SEGMENT_VERSIONS = [
    ("v5", "esg_segments_v5.jsonl"),
    ("v6", "esg_segments_v6.jsonl"),
    ("v7", "esg_segments_v7.jsonl"),
    ("v8", "esg_segments_v8.jsonl"),
]

CONFIG_FILE = "config.py"


def update_config(segment_file: str):
    """Update config.py with new segment file (version auto-increments)"""
    with open(CONFIG_FILE, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Update segment file path for esg_segments_v2 property
    import re
    content = re.sub(
        r'return os\.path\.join\(self\.processed_dir, "esg_segments_v\d+\.jsonl"\)',
        f'return os.path.join(self.processed_dir, "{segment_file}")',
        content
    )
    
    with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
        f.write(content)
    
    print(f"  Updated config to use: {segment_file}")


def run_inference(mode: str, llm_provider: str) -> tuple:
    """Run main.py with specified parameters, return (success, version)"""
    cmd = [sys.executable, "main.py", "--mode", mode, "--llm-provider", llm_provider]
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    # Extract version from output
    version = "unknown"
    for line in result.stdout.split('\n'):
        if "Version:" in line:
            import re
            match = re.search(r'v\d+', line)
            if match:
                version = match.group()
                break
    
    return result.returncode == 0, version


def main():
    parser = argparse.ArgumentParser(description="Batch test with different segment versions")
    parser.add_argument("--mode", choices=["train", "test"], default="test",
                        help="Inference mode (default: test)")
    parser.add_argument("--llm-provider", choices=["ollama", "gemini"], default="gemini",
                        help="LLM provider (default: gemini)")
    parser.add_argument("--segments", nargs="+", choices=["v5", "v6", "v7", "v8"], default=None,
                        help="Specify which segments to run, e.g., --segments v5 v7 (default: all)")
    args = parser.parse_args()
    
    # Filter segments if specified
    if args.segments:
        segments_to_run = [(name, file) for name, file in SEGMENT_VERSIONS if name in args.segments]
    else:
        segments_to_run = SEGMENT_VERSIONS
    
    print("=" * 60)
    print("🚀 Batch Test - Running segment files")
    print(f"   Mode: {args.mode}")
    print(f"   LLM Provider: {args.llm_provider}")
    print(f"   Segments: {[s[0] for s in segments_to_run]}")
    print(f"   Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("   Version: auto-increment based on existing results")
    print("=" * 60)
    
    results = []
    
    for seg_name, seg_file in segments_to_run:
        print(f"\n{'='*60}")
        print(f"📁 Running with {seg_name} ({seg_file})")
        print("=" * 60)
        
        # Update config (version auto-increments)
        update_config(seg_file)
        
        # Run inference
        success, version = run_inference(args.mode, args.llm_provider)
        results.append((seg_name, seg_file, version, success))
        
        if success:
            print(f"\n✅ {seg_name} ({version}) completed successfully")
        else:
            print(f"\n❌ {seg_name} ({version}) failed")
    
    # Summary
    print("\n" + "=" * 60)
    print("📊 BATCH TEST SUMMARY")
    print("=" * 60)
    for seg_name, seg_file, version, success in results:
        status = "✅ Success" if success else "❌ Failed"
        print(f"  {seg_name} ({version}): {status} - {seg_file}")
    
    # List output files
    print(f"\n📁 Output files in results/:")
    results_dir = "results"
    for seg_name, seg_file, version, _ in results:
        filename = f"{args.mode}_{version}_selective_merge.csv"
        filepath = os.path.join(results_dir, filename)
        if os.path.exists(filepath):
            size = os.path.getsize(filepath)
            print(f"  {filename} ({size} bytes)")
    
    print("\n✅ Batch test completed!")


if __name__ == "__main__":
    main()