| |
| """ |
| Expert Router for LEGO-MoE |
| Routes queries to appropriate experts based on ISL decisions |
| """ |
|
|
| import torch |
| import re |
| from pathlib import Path |
| import sys |
|
|
| sys.path.append(str(Path(__file__).parent / 'training' / 'distillation')) |
| from models import ClassificationExpert |
|
|
| class ExpertRouter: |
| """Routes queries to appropriate micro-experts""" |
| |
| def __init__(self): |
| print("Loading expert router...") |
| self.isl = self.load_isl_expert() |
| self.experts = self.load_all_experts() |
| print(f"✓ Router loaded with {len(self.experts)} experts") |
| |
| def load_isl_expert(self): |
| """Load ISL meta-governor""" |
| try: |
| expert = ClassificationExpert( |
| input_dim=256, |
| hidden_dim=256, |
| num_layers=3, |
| num_classes=3, |
| dropout=0.1 |
| ) |
| expert.load_state_dict(torch.load("models/isl_meta-governor_expert.pt", map_location='cpu')) |
| expert.eval() |
| return expert |
| except Exception as e: |
| print(f"⚠️ Could not load ISL expert: {e}") |
| return None |
| |
| def load_all_experts(self): |
| """Load all available micro-experts""" |
| experts = {} |
| |
| expert_configs = [ |
| ("grammar", 2), |
| ("mathematical_reasoning", 2), |
| ("question_answering", 2), |
| ("semantic_understanding", 2), |
| ] |
| |
| for name, num_classes in expert_configs: |
| try: |
| expert = ClassificationExpert( |
| input_dim=256, |
| hidden_dim=256, |
| num_layers=3, |
| num_classes=num_classes, |
| dropout=0.1 |
| ) |
| expert.load_state_dict(torch.load(f"models/{name}_expert.pt", map_location='cpu')) |
| expert.eval() |
| experts[name] = expert |
| except: |
| continue |
| |
| return experts |
| |
| def encode_query(self, query: str) -> torch.Tensor: |
| """Simple query encoding""" |
| features = torch.zeros(1, 256) |
| |
| |
| query_lower = query.lower() |
| |
| |
| if any(word in query_lower for word in ['calculate', 'compute', 'solve', 'math']): |
| features[0, 0] = 1.0 |
| if any(char in query for char in ['+', '-', '*', '/', '=', '%']): |
| features[0, 1] = 1.0 |
| if re.search(r'\d+', query): |
| features[0, 2] = 1.0 |
| |
| |
| if any(word in query_lower for word in ['grammar', 'correct', 'fix', 'spell']): |
| features[0, 3] = 1.0 |
| |
| |
| if query.strip().endswith('?'): |
| features[0, 4] = 1.0 |
| if any(word in query_lower for word in ['what', 'why', 'how', 'when', 'where', 'who']): |
| features[0, 5] = 1.0 |
| |
| |
| if any(word in query_lower for word in ['code', 'function', 'python', 'javascript']): |
| features[0, 6] = 1.0 |
| if any(char in query for char in ['(', ')', '{', '}', '[', ']', ';']): |
| features[0, 7] = 1.0 |
| |
| return features |
| |
| def isl_decision(self, query: str) -> str: |
| """ISL expert decides complexity""" |
| if not self.isl: |
| return "ESCALATE_FULL" |
| |
| embedding = self.encode_query(query) |
| |
| with torch.no_grad(): |
| _, confidence, logits = self.isl(embedding) |
| prediction = torch.argmax(logits, dim=1).item() |
| |
| decisions = ["TERMINATE_SAFE", "ESCALATE_FULL", "ESCALATE_PARTIAL"] |
| decision = decisions[prediction] if prediction < 3 else "ESCALATE_FULL" |
| |
| return decision |
| |
| def select_experts(self, query: str) -> list: |
| """Select which experts to run""" |
| selected = [] |
| query_lower = query.lower() |
| |
| |
| if any(word in query_lower for word in ['calculate', 'compute', 'solve', 'math', 'multiply', 'add', 'subtract']): |
| if "mathematical_reasoning" in self.experts: |
| selected.append("mathematical_reasoning") |
| |
| if re.search(r'\d+\s*[\+\-\*\/]\s*\d+', query): |
| if "mathematical_reasoning" in self.experts: |
| selected.append("mathematical_reasoning") |
| |
| |
| if any(word in query_lower for word in ['grammar', 'correct', 'fix', 'spell', 'punctuation']): |
| if "grammar" in self.experts: |
| selected.append("grammar") |
| |
| |
| if query.strip().endswith('?'): |
| if "question_answering" in self.experts: |
| selected.append("question_answering") |
| |
| |
| if len(query.split()) > 10: |
| if "semantic_understanding" in self.experts: |
| selected.append("semantic_understanding") |
| |
| return list(set(selected)) |
| |
| def route(self, query: str) -> dict: |
| """ |
| Route query to appropriate experts |
| |
| Returns: |
| { |
| 'decision': ISL decision, |
| 'experts': list of expert names to run, |
| 'direct_to_llm': bool |
| } |
| """ |
| |
| decision = self.isl_decision(query) |
| |
| |
| if decision == "TERMINATE_SAFE": |
| return { |
| 'decision': decision, |
| 'experts': [], |
| 'direct_to_llm': True |
| } |
| |
| |
| experts = self.select_experts(query) |
| |
| return { |
| 'decision': decision, |
| 'experts': experts, |
| 'direct_to_llm': len(experts) == 0 |
| } |
|
|
|
|
| def test_router(): |
| """Test expert router""" |
| print("="*70) |
| print("Testing Expert Router") |
| print("="*70) |
| |
| router = ExpertRouter() |
| |
| test_queries = [ |
| "What is 15 * 23?", |
| "Fix this sentence: I goes to school yesterday", |
| "What is the capital of France?", |
| "Hello, how are you?", |
| ] |
| |
| for query in test_queries: |
| print(f"\nQuery: {query}") |
| result = router.route(query) |
| print(f" ISL Decision: {result['decision']}") |
| print(f" Selected Experts: {result['experts']}") |
| print(f" Direct to LLM: {result['direct_to_llm']}") |
| |
| print("\n" + "="*70) |
| print("✓ Router working!") |
| print("="*70) |
|
|
|
|
| if __name__ == "__main__": |
| test_router() |
|
|
