| |
| """ |
| LEGO-MoE: Complete Hybrid MoE System |
| TinyLlama mother model + 23 micro-experts + ISL routing |
| """ |
|
|
| import sys |
| from pathlib import Path |
| from tinyllama_wrapper import TinyLlamaWrapper |
| from expert_router import ExpertRouter |
| import torch |
|
|
| sys.path.append(str(Path(__file__).parent / 'training' / 'distillation')) |
| from models import ClassificationExpert |
|
|
| class LegoMoE: |
| """Complete LEGO-MoE inference system""" |
| |
| def __init__(self): |
| print("="*70) |
| print("🧠 LEGO-MoE System Initialization") |
| print("="*70) |
| |
| |
| print("\n1. Loading TinyLlama mother model...") |
| self.tinyllama = TinyLlamaWrapper() |
| |
| print("\n2. Loading expert router...") |
| self.router = ExpertRouter() |
| |
| print("\n" + "="*70) |
| print("✅ LEGO-MoE System Ready!") |
| print("="*70) |
| |
| def run_expert(self, expert_name: str, query: str): |
| """Run a single expert on query""" |
| expert = self.router.experts.get(expert_name) |
| if not expert: |
| return f"Expert {expert_name} not available" |
| |
| |
| embedding = self.router.encode_query(query) |
| |
| |
| with torch.no_grad(): |
| _, confidence, logits = expert(embedding) |
| prediction = torch.argmax(logits, dim=1).item() |
| |
| |
| if expert_name == "mathematical_reasoning": |
| |
| return f"Prediction: {prediction}, Confidence: {confidence.item():.2f}" |
| elif expert_name == "grammar": |
| return "Correct" if prediction == 0 else "Needs correction" |
| else: |
| return f"Analysis: {prediction}" |
| |
| def generate(self, query: str, verbose: bool = True) -> str: |
| """ |
| Generate response using LEGO-MoE system |
| |
| Args: |
| query: User query |
| verbose: Show routing decisions |
| |
| Returns: |
| Generated response |
| """ |
| if verbose: |
| print(f"\n{'='*70}") |
| print(f"Query: {query}") |
| print(f"{'='*70}") |
| |
| |
| routing = self.router.route(query) |
| |
| if verbose: |
| print(f"\n🧭 ISL Decision: {routing['decision']}") |
| print(f"📋 Selected Experts: {routing['experts'] if routing['experts'] else 'None (direct to LLM)'}") |
| |
| |
| expert_outputs = {} |
| if routing['experts']: |
| if verbose: |
| print(f"\n🔬 Running {len(routing['experts'])} expert(s)...") |
| |
| for expert_name in routing['experts']: |
| output = self.run_expert(expert_name, query) |
| expert_outputs[expert_name] = output |
| |
| if verbose: |
| print(f" ✓ {expert_name}: {output}") |
| |
| |
| if verbose: |
| print(f"\n🤖 Generating response with TinyLlama...") |
| |
| if expert_outputs: |
| response = self.tinyllama.generate_with_context( |
| query, |
| expert_outputs, |
| temperature=0.7 |
| ) |
| else: |
| response = self.tinyllama.generate( |
| query, |
| temperature=0.7, |
| max_tokens=256 |
| ) |
| |
| if verbose: |
| print(f"\n{'='*70}") |
| print(f"Response:") |
| print(f"{'='*70}") |
| print(response) |
| print(f"{'='*70}\n") |
| |
| return response |
|
|
|
|
| def main(): |
| """CLI interface for LEGO-MoE""" |
| print("\n🧠 LEGO-MoE - Hybrid Mixture of Experts") |
| print("TinyLlama + 23 Micro-Experts + ISL Routing\n") |
| |
| |
| moe = LegoMoE() |
| |
| |
| if len(sys.argv) > 1: |
| |
| query = " ".join(sys.argv[1:]) |
| moe.generate(query) |
| else: |
| |
| print("\nInteractive Mode (type 'quit' to exit)") |
| print("="*70) |
| |
| while True: |
| try: |
| query = input("\nYou: ").strip() |
| |
| if not query: |
| continue |
| |
| if query.lower() in ['quit', 'exit', 'q']: |
| print("\nGoodbye! 👋") |
| break |
| |
| moe.generate(query) |
| |
| except KeyboardInterrupt: |
| print("\n\nGoodbye! 👋") |
| break |
| except Exception as e: |
| print(f"\nError: {e}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
