run_dual_probe.py

import os
import sys
from dual_probe_analysis import DualProbeAnalyzer

def main():
    
    # Configuration
    config = {
        'model_name': 'meta-llama/Llama-2-13b-chat-hf',
        'json_path': './data/storyanalogy_multiple_choice.json',
        'batch_size': 2,
        'device': 'cuda',
        'representation_type': 'hidden_states'
    }
    
    print(f"Model: {config['model_name']}")
    print(f"Dataset: {config['json_path']}")
    print(f"Representation type: {config['representation_type']}")
    print()
    
    # Check if dataset exists
    if not os.path.exists(config['json_path']):
        print(f"Dataset not found at: {config['json_path']}")
        print("Please check the path and try again.")
        return
    
    try:
        print("Initializing dual probe analyzer...")
        analyzer = DualProbeAnalyzer(config['model_name'])
        
        print("Running analogy probe analysis...")
        results, df_results = analyzer.run_dual_probe_analysis(
            config['json_path'], 
            config['batch_size'],
            config['representation_type']
        )
        
        print("\nAnalysis complete")
        
        # Show summary statistics
        analogy_cv_scores = [metrics['analogy_cv_mean'] for metrics in results.values()]
        print(f"\nAnalogy CV Mean Statistics:")
        print(f"  Mean: {sum(analogy_cv_scores)/len(analogy_cv_scores):.3f}")
        print(f"  Max:  {max(analogy_cv_scores):.3f}")
        print(f"  Min:  {min(analogy_cv_scores):.3f}")
        
    except Exception as e:
        print(f"Error during analysis: {str(e)}")
        print("Please check your configuration and try again.")
        import traceback
        traceback.print_exc()
        return

if __name__ == "__main__":
    main()