benchmark_ensemble.py - main - Paff/declarant

benchmark_ensemble.py
250 строк8 KB
1pash1985-gif
feat: AI Ensemble - DeepSeek + Kimi parallel parsing with debates
27 июл 2026, 11:14
27 июл 2026, 11:145ea7945
"""
Benchmark script for AI Ensemble debate parameters.

Tests different values of ENSEMBLE_DEBATE_MAX_ROUNDS and measures:
- Accuracy (compared to expected results)
- Execution time
- Number of disputes
- Confidence scores

Usage:
    python benchmark_ensemble.py [--rounds 0,1,2,3] [--sample 3] [--output benchmark_results.json]

Requires DEEPSEEK_API_KEY and KIMI_API_KEY in .env or environment.
"""

import argparse
import json
import os
import sys
import time
from pathlib import Path
from statistics import mean, stdev

# Add web/ to path
sys.path.insert(0, str(Path(__file__).parent / 'web'))


def load_sample_files(trip_folder: str, sample: int) -> list[tuple[str, str]]:
    """Load sample files from a trip folder.

    Returns list of (filename, text_content) tuples.
    """
    from ai_parser import extract_text_from_pdf, extract_text_from_excel

    trip_path = Path(trip_folder)
    if not trip_path.exists():
        print(f'Trip folder not found: {trip_folder}')
        return []

    files = []

    # Scan for PDF/Excel files
    for f in sorted(trip_path.rglob('*')):
        if f.suffix.lower() in ('.pdf', '.xls', '.xlsx'):
            files.append(f)
        elif f.suffix.lower() == '.rar':
            # Extract RAR archives
            try:
                import rarfile
                with rarfile.RarFile(f) as rf:
                    for info in rf.infolist():
                        if info.filename.lower().endswith(('.pdf', '.xls', '.xlsx')):
                            files.append(f'{f}::{info.filename}')
            except Exception as e:
                print(f'Warning: cannot read RAR {f}: {e}')

    # Limit sample
    if sample > 0:
        files = files[:sample]

    # Extract text
    result = []
    for f in files:
        try:
            if '::' in str(f):
                # RAR archive member
                rar_path, member = str(f).split('::', 1)
                import rarfile
                with rarfile.RarFile(rar_path) as rf:
                    data = rf.read(member)
                    # Save to temp file
                    temp_path = Path(__file__).parent / 'temp_extracted'
                    temp_path.mkdir(exist_ok=True)
                    temp_file = temp_path / Path(member).name
                    temp_file.write_bytes(data)
                    ext = temp_file.suffix.lower()
                    if ext == '.pdf':
                        text = extract_text_from_pdf(str(temp_file))
                    else:
                        text = extract_text_from_excel(str(temp_file))
                    temp_file.unlink()
            else:
                ext = f.suffix.lower()
                if ext == '.pdf':
                    text = extract_text_from_pdf(str(f))
                else:
                    text = extract_text_from_excel(str(f))

            if text and text.strip():
                result.append((str(f), text))
                print(f'  Loaded: {f.name} ({len(text)} chars)')
        except Exception as e:
            print(f'  Error loading {f}: {e}')

    return result


def run_benchmark(
    files: list[tuple[str, str]],
    rounds: int,
    timeout: int = 180,
) -> dict:
    """Run benchmark with specified debate rounds.

    Returns dict with metrics.
    """
    # Set environment variables
    os.environ['ENSEMBLE_DEBATE_MAX_ROUNDS'] = str(rounds)
    os.environ['ENSEMBLE_TIMEOUT_SECONDS'] = str(timeout)

    # Reload config to pick up new values
    import importlib
    import config
    importlib.reload(config)

    from ai_ensemble import parse_invoice_with_ensemble

    results = []
    total_time = 0
    total_disputes = 0
    total_confidence = 0

    print(f'\n=== Benchmark: rounds={rounds} ===')

    for filename, text in files:
        print(f'  Parsing: {Path(filename).name}...', end='', flush=True)

        start = time.time()
        try:
            result = parse_invoice_with_ensemble(text)
            elapsed = time.time() - start

            # Extract metrics
            meta = result.get('_ensemble_meta', {})
            disputes = len(meta.get('disputes', []))
            confidence = meta.get('confidence', 0)
            rounds_used = meta.get('rounds', 0)

            results.append({
                'file': Path(filename).name,
                'time': elapsed,
                'disputes': disputes,
                'confidence': confidence,
                'rounds_used': rounds_used,
                'supplier': result.get('supplier', ''),
                'items_count': len(result.get('items', [])),
                'error': result.get('error'),
            })

            total_time += elapsed
            total_disputes += disputes
            total_confidence += confidence

            status = 'OK' if 'error' not in result else f"ERROR: {result.get('error')}"
            print(f' {elapsed:.1f}s, {disputes} disputes, conf={confidence:.2f} [{status}]')

        except Exception as e:
            elapsed = time.time() - start
            results.append({
                'file': Path(filename).name,
                'time': elapsed,
                'error': str(e),
            })
            print(f' ERROR: {e}')

    # Aggregate metrics
    successful = [r for r in results if 'error' not in r or not r.get('error')]
    metrics = {
        'rounds': rounds,
        'files_count': len(files),
        'successful': len(successful),
        'total_time': total_time,
        'avg_time': mean([r['time'] for r in successful]) if successful else 0,
        'total_disputes': total_disputes,
        'avg_disputes': mean([r['disputes'] for r in successful]) if successful else 0,
        'avg_confidence': mean([r['confidence'] for r in successful]) if successful else 0,
        'results': results,
    }

    return metrics


def main():
    parser = argparse.ArgumentParser(description='AI Ensemble benchmark')
    parser.add_argument('--trip', default=r'c:\Users\0\Desktop\Задача по логистике\26.07\26.07\BOJ444',
                        help='Trip folder with sample files')
    parser.add_argument('--sample', type=int, default=3,
                        help='Number of files to test (0 = all)')
    parser.add_argument('--rounds', default='0,1,2',
                        help='Comma-separated debate rounds to test')
    parser.add_argument('--timeout', type=int, default=180,
                        help='Timeout per file in seconds')
    parser.add_argument('--output', default='benchmark_results.json',
                        help='Output JSON file')

    args = parser.parse_args()

    print('AI Ensemble Benchmark')
    print('=' * 60)
    print(f'Trip: {args.trip}')
    print(f'Sample: {args.sample} files')
    print(f'Rounds to test: {args.rounds}')
    print(f'Timeout: {args.timeout}s')
    print()

    # Load sample files
    print('Loading sample files...')
    files = load_sample_files(args.trip, args.sample)
    if not files:
        print('No files loaded. Exiting.')
        return

    print(f'\nLoaded {len(files)} files.')

    # Parse rounds
    rounds_list = [int(r.strip()) for r in args.rounds.split(',')]

    # Run benchmarks
    all_metrics = []
    for rounds in rounds_list:
        metrics = run_benchmark(files, rounds, args.timeout)
        all_metrics.append(metrics)

    # Summary
    print('\n' + '=' * 60)
    print('SUMMARY')
    print('=' * 60)
    print(f'{"Rounds":<10} {"Avg Time":<12} {"Disputes":<12} {"Confidence":<12}')
    print('-' * 60)

    for m in all_metrics:
        print(f'{m["rounds"]:<10} {m["avg_time"]:<12.1f} {m["avg_disputes"]:<12.1f} {m["avg_confidence"]:<12.2f}')

    # Save results
    output_path = Path(args.output)
    output_path.write_text(
        json.dumps(all_metrics, ensure_ascii=False, indent=2),
        encoding='utf-8'
    )
    print(f'\nResults saved to: {output_path}')

    # Recommendation
    print('\nRECOMMENDATION:')
    best = min(all_metrics, key=lambda m: m['avg_time'] * (1 + m['avg_disputes'] * 0.1))
    print(f'  Optimal rounds: {best["rounds"]}')
    print(f'  Reason: Best balance of speed ({best["avg_time"]:.1f}s) and quality ({best["avg_disputes"]:.1f} disputes)')


if __name__ == '__main__':
    main()