run_bird_experiments.py

#!/usr/bin/env python3
"""

"""

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

# 
sys.path.append(str(Path(__file__).parent))

from experiments.bird_experiment import BIRDExperiment, run_comparison_experiment
from src.utils.llm_client import LLMManager


def check_environment():
    """"""
    print("="*70)
    print("Environment Check")
    print("="*70)
    
    #  API Keys
    manager = LLMManager()
    available_models = manager.get_available_models()
    
    print(f"\nAvailable Models: {available_models}")
    
    if not available_models:
        print("\nWARNING  Warning: No LLM models available!")
        print("Please set the following environment variables in .env file:")
        print("  - DEEPSEEK_API_KEY")
        print("  - QWEN_API_KEY")
        return False
    
    print(f"\nOK {len(available_models)} model(s) ready")
    return True


def run_basic_test():
    """"""
    print("\n" + "="*70)
    print("Running Basic Tests")
    print("="*70)
    
    # 
    print("\n1. Testing data loader...")
    from src.data.bird_loader import BIRDDataset
    dataset = BIRDDataset()
    dev_data = dataset.load_dev_data()
    print(f"   OK Loaded {len(dev_data)} examples")
    
    #  LLM
    print("\n2. Testing LLM client...")
    manager = LLMManager()
    if manager.get_available_models():
        model = manager.get_available_models()[0]
        test_prompt = "Generate a simple SQL query to select all columns from a table named 'users'."
        response = manager.generate(test_prompt, model=model, max_tokens=100)
        print(f"   OK {model} responded: {response[:100]}...")
    else:
        print("   ERROR No LLM available")
    
    # 
    print("\n3. Testing evaluator...")
    from src.utils.evaluator import BIRDEvaluator
    evaluator = BIRDEvaluator()
    print("   OK Evaluator initialized")
    
    print("\nOK All basic tests passed!")


def run_quick_experiment():
    """"""
    print("\n" + "="*70)
    print("Running Quick Experiment (5 samples)")
    print("="*70)
    
    manager = LLMManager()
    available_models = manager.get_available_models()
    
    if not available_models:
        print("No models available, skipping experiment")
        return
    
    model = available_models[0]
    
    experiment = BIRDExperiment(
        model=model,
        use_schema_info=True,
        use_evidence=True
    )
    
    predictions = experiment.run(
        num_samples=5,
        output_file=f"quick_test_{model}.json"
    )
    
    metrics = experiment.evaluate(predictions)
    
    return {
        "model": model,
        "samples": 5,
        "metrics": metrics
    }


def generate_report(results: dict):
    """"""
    report_time = datetime.now().strftime("%Y%m%d_%H%M%S")
    report_file = f"experiments/results/report_{report_time}.md"
    
    Path("experiments/results").mkdir(parents=True, exist_ok=True)
    
    with open(report_file, 'w') as f:
        f.write("# BIRD Benchmark Experiment Report\n\n")
        f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
        
        f.write("## Environment\n\n")
        manager = LLMManager()
        f.write(f"- Available Models: {', '.join(manager.get_available_models())}\n")
        f.write(f"- Total Examples: {results.get('total_examples', 'N/A')}\n\n")
        
        if 'quick_experiment' in results:
            f.write("## Quick Experiment Results\n\n")
            qe = results['quick_experiment']
            f.write(f"- Model: {qe['model']}\n")
            f.write(f"- Samples: {qe['samples']}\n")
            metrics = qe['metrics']
            f.write(f"- Exact Match Rate: {metrics.get('exact_match_rate', 0):.2f}%\n")
            f.write(f"- Execution Accuracy: {metrics.get('execution_accuracy_rate', 0):.2f}%\n\n")
    
    print(f"\nOK Report saved to {report_file}")


def main():
    """"""
    print("\n" + "="*70)
    print("Multimodal Ambiguity Resolution BIRD Benchmark Experiment Runner")
    print("="*70)
    
    results = {}
    
    # 1. 
    if not check_environment():
        print("\nWARNING  Environment check failed. Please configure API keys in .env file.")
        return
    
    # 2. 
    try:
        run_basic_test()
    except Exception as e:
        print(f"\nERROR Basic test failed: {e}")
        return
    
    # 3. 
    try:
        quick_result = run_quick_experiment()
        if quick_result:
            results['quick_experiment'] = quick_result
            results['total_examples'] = quick_result['samples']
    except Exception as e:
        print(f"\nERROR Quick experiment failed: {e}")
    
    # 4. 
    generate_report(results)
    
    print("\n" + "="*70)
    print("Experiment Runner Completed!")
    print("="*70)
    print("\nNext steps:")
    print("  1. Check results in experiments/results/")
    print("  2. Run full experiment: python experiments/bird_experiment.py --samples 100")
    print("  3. Run comparison: python experiments/bird_experiment.py --compare")


if __name__ == "__main__":
    main()