File size: 4,788 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 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | #!/usr/bin/env python3
"""
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)
# Load components
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"
# Encode query
embedding = self.router.encode_query(query)
# Run expert
with torch.no_grad():
_, confidence, logits = expert(embedding)
prediction = torch.argmax(logits, dim=1).item()
# Format output based on expert type
if expert_name == "mathematical_reasoning":
# For math, try to extract answer
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}")
# 1. Route with ISL
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)'}")
# 2. Run experts if needed
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}")
# 3. Generate with TinyLlama
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")
# Initialize system
moe = LegoMoE()
# Interactive mode
if len(sys.argv) > 1:
# Single query mode
query = " ".join(sys.argv[1:])
moe.generate(query)
else:
# Interactive mode
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()
|