File size: 5,002 Bytes
ae1e90a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | #!/usr/bin/env python3
"""
Enhanced Expert Router - Loads All 23 Trained Experts
"""
import torch
from pathlib import Path
import sys
sys.path.append(str(Path(__file__).parent / 'training' / 'distillation'))
from models import ClassificationExpert
class EnhancedExpertRouter:
"""Enhanced router with all 23 experts loaded"""
def __init__(self):
print("Loading Enhanced Expert Router...")
self.experts = {}
self.expert_configs = self._get_all_expert_configs()
self._load_all_experts()
print(f"✓ Loaded {len(self.experts)} experts")
def _get_all_expert_configs(self):
"""Get configurations for all 23 experts"""
return {
# General Experts (15)
"grammar": {"num_classes": 2, "category": "general"},
"grammar_v2": {"num_classes": 2, "category": "general"},
"mathematical_reasoning": {"num_classes": 2, "category": "general"},
"logical_inference": {"num_classes": 2, "category": "general"},
"common_sense": {"num_classes": 2, "category": "general"},
"question_answering": {"num_classes": 2, "category": "general"},
"code_generation": {"num_classes": 2, "category": "general"},
"sentiment_analysis": {"num_classes": 3, "category": "general"},
"toxicity_detection": {"num_classes": 2, "category": "general"},
"semantic_understanding": {"num_classes": 2, "category": "general"},
"summarization": {"num_classes": 2, "category": "general"},
"ner": {"num_classes": 9, "category": "general"},
"translation": {"num_classes": 2, "category": "general"},
"causal_reasoning": {"num_classes": 2, "category": "general"},
# Security Experts (5)
"attack_surface_mapper": {"num_classes": 2, "category": "security"},
"authorization_logic": {"num_classes": 2, "category": "security"},
"data_flow_analysis": {"num_classes": 2, "category": "security"},
"business_logic": {"num_classes": 2, "category": "security"},
"exploit_feasibility": {"num_classes": 2, "category": "security"},
# Meta-Cognitive (1)
"isl_meta-governor": {"num_classes": 3, "category": "meta"},
}
def _load_all_experts(self):
"""Load all available expert models"""
models_dir = Path("models")
for expert_name, config in self.expert_configs.items():
model_path = models_dir / f"{expert_name}_expert.pt"
if not model_path.exists():
print(f" ⚠️ {expert_name}: model file not found")
continue
try:
expert = ClassificationExpert(
input_dim=256,
hidden_dim=256,
num_layers=3,
num_classes=config["num_classes"],
dropout=0.1
)
expert.load_state_dict(torch.load(model_path, map_location='cpu'))
expert.eval()
self.experts[expert_name] = {
"model": expert,
"category": config["category"],
"num_classes": config["num_classes"]
}
print(f" ✓ {expert_name} ({config['category']})")
except Exception as e:
print(f" ✗ {expert_name}: {str(e)[:50]}")
def get_expert(self, name: str):
"""Get expert by name"""
return self.experts.get(name)
def get_experts_by_category(self, category: str):
"""Get all experts in a category"""
return {
name: info for name, info in self.experts.items()
if info["category"] == category
}
def get_summary(self):
"""Get summary of loaded experts"""
by_category = {}
for name, info in self.experts.items():
cat = info["category"]
if cat not in by_category:
by_category[cat] = []
by_category[cat].append(name)
return {
"total": len(self.experts),
"by_category": by_category
}
def test_enhanced_router():
"""Test enhanced router"""
print("="*70)
print("Testing Enhanced Expert Router")
print("="*70)
router = EnhancedExpertRouter()
print("\n" + "="*70)
print("Summary:")
summary = router.get_summary()
print(f"Total experts loaded: {summary['total']}")
for category, experts in summary['by_category'].items():
print(f"\n{category.upper()} ({len(experts)}):")
for expert in sorted(experts):
print(f" - {expert}")
print("\n" + "="*70)
print("✓ Enhanced router working!")
print("="*70)
if __name__ == "__main__":
test_enhanced_router()
|