biomind_diagnostic_fixed.py

#!/usr/bin/env python3
"""
BIOMIND System Diagnostic - Fixed Version
Tests the actual BIOMIND components that are working
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

import torch
import torch.nn.functional as F
import numpy as np
import time
import json
from typing import Dict, Any, List
from datetime import datetime

def test_integrated_biomind_evaluation():
    """Test the integrated BIOMIND system on MMLU samples"""
    print("[BRAIN] Testing Integrated BIOMIND Evaluation System...")
    
    try:
        from integrated_biomind_evaluation import IntegratedBIOMINDEvaluator
        
        # Initialize the evaluator
        specialists = ['qwen_math_expert', 'qwen_general_reasoner', 'tiny_llama_planner', 'tiny_llama_critic']
        evaluator = IntegratedBIOMINDEvaluator(specialists)
        print("   [OK] IntegratedBIOMINDEvaluator initialized")
        
        # Test MMLU evaluation
        print("   [BRAIN] Testing MMLU evaluation...")
        
        # Run a small MMLU test
        from real_dataset_loader import RealDatasetLoader
        
        # Load MMLU samples
        loader = RealDatasetLoader()
        mmlu_samples = loader.load_mmlu_full(max_samples=10)
        
        print(f"      Loaded {len(mmlu_samples)} MMLU samples")
        
        # Evaluate samples
        correct = 0
        total = 0
        processing_times = []
        
        for i, sample in enumerate(mmlu_samples):
            start_time = time.time()
            
            question = sample['question']
            choices = sample['choices']
            correct_answer = sample['answer']
            subject = sample.get('subject', 'general')
            
            # Evaluate with integrated BIOMIND
            result = evaluator.evaluate_question(
                question=question,
                choices=choices,
                correct_answer=correct_answer,
                subject=subject
            )
            
            processing_time = (time.time() - start_time) * 1000
            processing_times.append(processing_time)
            
            if result['is_correct']:
                correct += 1
            total += 1
            
            print(f"      Sample {i+1}: {'[OK]' if result['is_correct'] else '[FAIL]'} "
                  f"({result['selected_specialist']}) {processing_time:.1f}ms")
        
        accuracy = correct / total if total > 0 else 0
        avg_time = np.mean(processing_times) if processing_times else 0
        
        print(f"   [CHART] MMLU Test Results:")
        print(f"      Accuracy: {accuracy:.1%} ({correct}/{total})")
        print(f"      Avg Time: {avg_time:.1f}ms")
        
        return {
            'component': 'integrated_biomind',
            'test': 'mmlu_evaluation',
            'accuracy': accuracy,
            'correct': correct,
            'total': total,
            'avg_time_ms': avg_time,
            'passed': accuracy > 0.5  # Pass if > 50%
        }
        
    except Exception as e:
        print(f"   [FAIL] Integrated BIOMIND test failed: {e}")
        import traceback
        traceback.print_exc()
        return {
            'component': 'integrated_biomind',
            'test': 'mmlu_evaluation', 
            'passed': False,
            'error': str(e)
        }

def test_true_biomind_evaluation():
    """Test the TRUE BIOMIND system on MMLU samples"""
    print("[BRAIN] Testing TRUE BIOMIND Evaluation System...")
    
    try:
        from true_biomind_evaluation import TrueBIOMINDEvaluator
        
        # Initialize the evaluator
        evaluator = TrueBIOMINDEvaluator()
        print("   [OK] TrueBIOMINDEvaluator initialized")
        
        # Test MMLU evaluation
        print("   [BRAIN] Testing MMLU evaluation...")
        
        # Load MMLU samples
        from real_dataset_loader import RealDatasetLoader
        loader = RealDatasetLoader()
        mmlu_samples = loader.load_mmlu_full(max_samples=10)
        
        print(f"      Loaded {len(mmlu_samples)} MMLU samples")
        
        # Evaluate samples
        correct = 0
        total = 0
        processing_times = []
        
        for i, sample in enumerate(mmlu_samples):
            start_time = time.time()
            
            question = sample['question']
            choices = sample['choices']
            correct_answer = sample['answer']
            subject = sample.get('subject', 'general')
            
            # Evaluate with TRUE BIOMIND
            result = evaluator._evaluate_with_real_biomind_architecture(
                question=question,
                choices=choices,
                correct_answer=correct_answer,
                subject=subject
            )
            
            processing_time = (time.time() - start_time) * 1000
            processing_times.append(processing_time)
            
            if result['is_correct']:
                correct += 1
            total += 1
            
            print(f"      Sample {i+1}: {'[OK]' if result['is_correct'] else '[FAIL]'} "
                  f"({result['selected_specialist']}) {processing_time:.1f}ms")
        
        accuracy = correct / total if total > 0 else 0
        avg_time = np.mean(processing_times) if processing_times else 0
        
        print(f"   [CHART] TRUE BIOMIND Test Results:")
        print(f"      Accuracy: {accuracy:.1%} ({correct}/{total})")
        print(f"      Avg Time: {avg_time:.1f}ms")
        
        return {
            'component': 'true_biomind',
            'test': 'mmlu_evaluation',
            'accuracy': accuracy,
            'correct': correct,
            'total': total,
            'avg_time_ms': avg_time,
            'passed': accuracy > 0.1  # More lenient for TRUE system
        }
        
    except Exception as e:
        print(f"   [FAIL] TRUE BIOMIND test failed: {e}")
        import traceback
        traceback.print_exc()
        return {
            'component': 'true_biomind',
            'test': 'mmlu_evaluation',
            'passed': False,
            'error': str(e)
        }

def test_thalamic_router_standalone():
    """Test the thalamic router component independently"""
    print("[BRAIN] Testing Standalone Thalamic Router...")
    
    try:
        # Test the simple thalamic router training
        from train_simple_thalamic import SimpleTrainer, SimpleThalamicRouter
        
        # Initialize trainer
        specialists = ['qwen_math_expert', 'qwen_general_reasoner', 'tiny_llama_planner', 'tiny_llama_critic']
        trainer = SimpleTrainer(specialists=specialists)
        print("   [OK] SimpleTrainer initialized")
        
        # Train for a few epochs
        print("   ? Training thalamic router...")
        trainer.train_epochs(epochs=2, examples_per_epoch=100)
        
        # Test routing decisions
        test_questions = [
            "What is the derivative of x²?",
            "What is the primary function of mitochondria?",
            "What is the best strategy for project planning?",
            "Evaluate the complexity of binary search"
        ]
        
        correct_predictions = 0
        total_predictions = len(test_questions)
        processing_times = []
        
        for i, question in enumerate(test_questions):
            start_time = time.time()
            
            trainer.router.eval()
            with torch.no_grad():
                result = trainer.router(question)
            
            processing_time = (time.time() - start_time) * 1000
            processing_times.append(processing_time)
            
            predicted_specialist = result['predicted_specialist']
            confidence = result['confidence']
            
            # Check if routing makes sense
            routing_makes_sense = False
            if "derivative" in question.lower() and "math" in predicted_specialist:
                routing_makes_sense = True
            elif "mitochondria" in question.lower() and "reasoner" in predicted_specialist:
                routing_makes_sense = True
            elif "strategy" in question.lower() and "planner" in predicted_specialist:
                routing_makes_sense = True
            elif "complexity" in question.lower() and "critic" in predicted_specialist:
                routing_makes_sense = True
            elif confidence > 0.2:  # Any reasonable confidence is ok
                routing_makes_sense = True
                
            if routing_makes_sense:
                correct_predictions += 1
                
            print(f"      Q{i+1}: {predicted_specialist} (conf:{confidence:.2f}) {processing_time:.1f}ms "
                  f"{'[OK]' if routing_makes_sense else '[FAIL]'}")
        
        accuracy = correct_predictions / total_predictions
        avg_time = np.mean(processing_times)
        
        print(f"   [CHART] Thalamic Router Results:")
        print(f"      Routing Accuracy: {accuracy:.1%} ({correct_predictions}/{total_predictions})")
        print(f"      Avg Time: {avg_time:.1f}ms")
        
        return {
            'component': 'thalamic_router',
            'test': 'routing_accuracy',
            'accuracy': accuracy,
            'correct': correct_predictions,
            'total': total_predictions,
            'avg_time_ms': avg_time,
            'passed': accuracy > 0.5
        }
        
    except Exception as e:
        print(f"   [FAIL] Thalamic router test failed: {e}")
        import traceback
        traceback.print_exc()
        return {
            'component': 'thalamic_router',
            'test': 'routing_accuracy',
            'passed': False,
            'error': str(e)
        }

def run_comprehensive_diagnostic():
    """Run comprehensive BIOMIND diagnostic"""
    print("[ROCKET] BIOMIND System Comprehensive Diagnostic")
    print("=" * 60)
    
    start_time = time.time()
    
    # Run all tests
    results = []
    
    # Test 1: Integrated BIOMIND (the 100% system)
    integrated_result = test_integrated_biomind_evaluation()
    results.append(integrated_result)
    
    print()
    
    # Test 2: TRUE BIOMIND (the real system)
    true_result = test_true_biomind_evaluation()
    results.append(true_result)
    
    print()
    
    # Test 3: Standalone Thalamic Router
    router_result = test_thalamic_router_standalone()
    results.append(router_result)
    
    total_time = time.time() - start_time
    
    # Calculate overall metrics
    total_tests = len(results)
    passed_tests = sum(1 for r in results if r.get('passed', False))
    overall_success = passed_tests / total_tests if total_tests > 0 else 0
    
    # Generate summary
    print("\n" + "=" * 60)
    print("[MDN] DIAGNOSTIC SUMMARY")
    print("=" * 60)
    
    print(f"Overall Success Rate: {overall_success:.1%}")
    print(f"Tests Passed: {passed_tests}/{total_tests}")
    print(f"Total Diagnostic Time: {total_time:.1f}s")
    
    print(f"\n[CHART] Component Results:")
    for result in results:
        component = result['component']
        test_name = result['test']
        passed = result.get('passed', False)
        status = "[OK] PASS" if passed else "[FAIL] FAIL"
        
        print(f"  ? {component} ({test_name}): {status}")
        
        if 'accuracy' in result:
            print(f"    - Accuracy: {result['accuracy']:.1%}")
        if 'avg_time_ms' in result:
            print(f"    - Avg Time: {result['avg_time_ms']:.1f}ms")
        if 'error' in result:
            print(f"    - Error: {result['error']}")
    
    # Performance comparison
    integrated_acc = None
    true_acc = None
    
    for result in results:
        if result['component'] == 'integrated_biomind' and 'accuracy' in result:
            integrated_acc = result['accuracy']
        elif result['component'] == 'true_biomind' and 'accuracy' in result:
            true_acc = result['accuracy']
    
    if integrated_acc is not None and true_acc is not None:
        print(f"\n[SEARCH] PERFORMANCE ANALYSIS:")
        print(f"  ? Integrated BIOMIND: {integrated_acc:.1%}")
        print(f"  ? TRUE BIOMIND: {true_acc:.1%}")
        print(f"  ? Performance Gap: {(integrated_acc - true_acc)*100:.1f} percentage points")
        
        if integrated_acc > 0.8 and true_acc < 0.2:
            print("  [WARN]  MAJOR DISCREPANCY DETECTED!")
            print("  ? Integrated system likely using simulation")
            print("  ? TRUE system shows realistic AI performance")
    
    # Save results
    diagnostic_data = {
        'timestamp': datetime.now().isoformat(),
        'overall_success_rate': overall_success,
        'total_tests': total_tests,
        'passed_tests': passed_tests,
        'total_time_s': total_time,
        'results': results
    }
    
    output_file = f"biomind_diagnostic_results_{int(time.time())}.json"
    with open(output_file, 'w') as f:
        json.dump(diagnostic_data, f, indent=2, default=str)
    
    print(f"\n? Detailed results saved to: {output_file}")
    
    return diagnostic_data

def main():
    """Main diagnostic function"""
    try:
        results = run_comprehensive_diagnostic()
        return results
    except Exception as e:
        print(f"[FAIL] Diagnostic failed: {e}")
        import traceback
        traceback.print_exc()
        return None

if __name__ == "__main__":
    main()