| |
| """ |
| 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 { |
| |
| "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"}, |
| |
| |
| "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"}, |
| |
| |
| "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() |
|
|
