beastbullet-experts / code /hybrid_inference.py

Complete package: 18 experts + full codebase

ae1e90a 11 months ago

9.78 kB

#!/usr/bin/env python3
"""
Hybrid Inference: OLMoE-0125 + 18 Micro-Experts
Phase 1: Hint Injection for immediate quality boost
"""

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from pathlib import Path
import json
import time
import sys

# Add paths
sys.path.append(str(Path(__file__).parent / 'training' / 'encoder'))
sys.path.append(str(Path(__file__).parent / 'training' / 'distillation'))

from model import create_encoder
from models import ClassificationExpert

class HybridMoESystem:
    """Hierarchical Hybrid MoE: OLMoE Mother + Micro-Experts"""
    
    def __init__(
        self,
        olmoe_model_path: str,
        encoder_path: Path,
        experts_dir: Path,
        device: str = "auto"
    ):
        self.device = device
        
        # Load OLMoE Mother
        print("🔄 Loading OLMoE-0125 Mother...")
        self.olmoe = AutoModelForCausalLM.from_pretrained(
            olmoe_model_path,
            trust_remote_code=True,
            torch_dtype=torch.float16,
            device_map=device
        )
        self.tokenizer = AutoTokenizer.from_pretrained(olmoe_model_path)
        print("✓ OLMoE loaded")
        
        # Load Shared Encoder
        print("\n🔄 Loading Shared Encoder...")
        self.encoder_tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
        self.encoder = create_encoder({
            'vocab_size': self.encoder_tokenizer.vocab_size,
            'embed_dim': 256,
            'num_layers': 4,
            'num_heads': 8,
            'ff_dim': 1024,
            'dropout': 0.1,
            'max_seq_len': 128
        })
        self.encoder.load_state_dict(torch.load(encoder_path, map_location='cpu'))
        self.encoder.eval()
        print("✓ Encoder loaded")
        
        # Load Micro-Experts
        print("\n🔄 Loading Micro-Experts...")
        self.experts = self._load_experts(experts_dir)
        print(f"✓ {len(self.experts)} experts loaded")
    
    def _load_experts(self, experts_dir: Path):
        """Load all trained micro-experts"""
        experts = {}
        
        # Grammar Expert
        grammar_path = experts_dir / 'grammar_expert_v2.pt'
        if grammar_path.exists():
            expert = ClassificationExpert(256, 256, 2, 3, 0.1)
            expert.load_state_dict(torch.load(grammar_path, map_location='cpu'))
            expert.eval()
            experts['grammar'] = {
                'model': expert,
                'name': 'Grammar & Syntax',
                'threshold': 0.7,
                'hint_template': "[GRAMMAR EXPERT]: {advice}"
            }
        
        # Math Expert
        math_path = experts_dir / 'mathematical_reasoning_expert.pt'
        if math_path.exists():
            expert = ClassificationExpert(256, 256, 1, 3, 0.1)
            expert.load_state_dict(torch.load(math_path, map_location='cpu'))
            expert.eval()
            experts['math'] = {
                'model': expert,
                'name': 'Mathematical Reasoning',
                'threshold': 0.8,
                'hint_template': "[MATH EXPERT]: {advice}"
            }
        
        # Sentiment Expert
        sentiment_path = experts_dir / 'sentiment_analysis_expert.pt'
        if sentiment_path.exists():
            expert = ClassificationExpert(256, 256, 4, 3, 0.1)
            expert.load_state_dict(torch.load(sentiment_path, map_location='cpu'))
            expert.eval()
            experts['sentiment'] = {
                'model': expert,
                'name': 'Sentiment Analysis',
                'threshold': 0.7,
                'hint_template': "[SENTIMENT EXPERT]: {advice}"
            }
        
        return experts
    
    def encode_prompt(self, prompt: str):
        """Encode prompt with shared encoder"""
        inputs = self.encoder_tokenizer(
            prompt,
            padding='max_length',
            truncation=True,
            max_length=128,
            return_tensors='pt'
        )
        
        with torch.no_grad():
            embedding = self.encoder(inputs['input_ids'], inputs['attention_mask'])
        
        return embedding
    
    def run_experts(self, embedding):
        """Run all experts in parallel"""
        results = []
        
        for expert_id, expert_data in self.experts.items():
            model = expert_data['model']
            
            with torch.no_grad():
                output_emb, confidence, logits = model(embedding)
            
            # Get confidence score
            probs = torch.softmax(logits, dim=-1)
            max_conf = probs.max().item()
            predicted_class = probs.argmax().item()
            
            results.append({
                'id': expert_id,
                'name': expert_data['name'],
                'confidence': max_conf,
                'prediction': predicted_class,
                'threshold': expert_data['threshold'],
                'template': expert_data['hint_template']
            })
        
        return results
    
    def generate_hints(self, expert_results, prompt):
        """Generate expert hints for high-confidence activations"""
        hints = []
        
        for result in expert_results:
            if result['confidence'] > result['threshold']:
                # Generate domain-specific advice
                if result['id'] == 'grammar':
                    advice = "Ensure perfect grammar and sentence structure."
                elif result['id'] == 'math':
                    advice = "Provide step-by-step mathematical reasoning."
                elif result['id'] == 'sentiment':
                    advice = "Consider emotional tone and sentiment."
                else:
                    advice = f"Apply {result['name']} expertise."
                
                hint = result['template'].format(advice=advice)
                hints.append(hint)
        
        return hints
    
    def hybrid_generate(
        self,
        prompt: str,
        max_new_tokens: int = 200,
        temperature: float = 0.7,
        top_k: int = 4
    ):
        """Hybrid generation with expert hint injection"""
        
        print(f"\n{'='*70}")
        print(f"🎯 Hybrid Inference")
        print(f"{'='*70}")
        print(f"Prompt: {prompt}\n")
        
        # Step 1: Encode prompt
        print("1️⃣ Encoding prompt...")
        embedding = self.encode_prompt(prompt)
        
        # Step 2: Run experts
        print("2️⃣ Running micro-experts...")
        expert_results = self.run_experts(embedding)
        
        # Step 3: Filter by confidence
        active_experts = [
            e for e in expert_results 
            if e['confidence'] > e['threshold']
        ]
        
        print(f"   Active experts: {len(active_experts)}/{len(expert_results)}")
        for exp in active_experts:
            print(f"   - {exp['name']}: {exp['confidence']:.2%}")
        
        # Step 4: Select top-k
        top_experts = sorted(
            active_experts,
            key=lambda x: x['confidence'],
            reverse=True
        )[:top_k]
        
        # Step 5: Generate hints
        print("\n3️⃣ Generating expert hints...")
        hints = self.generate_hints(top_experts, prompt)
        
        # Step 6: Inject into system prompt
        if hints:
            hint_text = "\n".join(hints)
            enriched_prompt = f"""{hint_text}
User: {prompt}
Assistant:"""
        else:
            enriched_prompt = f"User: {prompt}\nAssistant:"
        
        print(f"\n4️⃣ OLMoE generation...")
        
        # Step 7: Generate with OLMoE
        inputs = self.tokenizer(enriched_prompt, return_tensors="pt").to(self.olmoe.device)
        
        start = time.time()
        outputs = self.olmoe.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=True
        )
        elapsed = time.time() - start
        
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract assistant response
        if "Assistant:" in response:
            response = response.split("Assistant:")[-1].strip()
        
        tokens_generated = len(outputs[0]) - len(inputs['input_ids'][0])
        tokens_per_sec = tokens_generated / elapsed
        
        print(f"\n{'='*70}")
        print(f"📊 Results")
        print(f"{'='*70}")
        print(f"Response: {response}\n")
        print(f"⚡ Performance:")
        print(f"   Tokens: {tokens_generated}")
        print(f"   Time: {elapsed:.2f}s")
        print(f"   Speed: {tokens_per_sec:.1f} t/s")
        print(f"   Active Experts: {len(top_experts)}")
        print(f"{'='*70}\n")
        
        return {
            'response': response,
            'active_experts': top_experts,
            'tokens_per_sec': tokens_per_sec,
            'hints': hints
        }


def main():
    print("="*70)
    print("🚀 Hybrid MoE Inference System")
    print("   OLMoE-0125 + Micro-Experts")
    print("="*70)
    
    # Initialize system
    system = HybridMoESystem(
        olmoe_model_path="models/olmoe-0125",
        encoder_path=Path("models/encoder_bootstrap.pt"),
        experts_dir=Path("models"),
        device="auto"
    )
    
    # Test prompts
    test_prompts = [
        "What is 15 + 27?",
        "Fix this sentence: Me and him goes to the store yesterday",
        "Write a Python function to calculate factorial",
        "This movie was absolutely terrible and I hated it",
    ]
    
    for prompt in test_prompts:
        result = system.hybrid_generate(prompt)
        time.sleep(1)  # Brief pause between tests


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""
	Hybrid Inference: OLMoE-0125 + 18 Micro-Experts
	Phase 1: Hint Injection for immediate quality boost
	"""

	import torch
	import torch.nn as nn
	from transformers import AutoModelForCausalLM, AutoTokenizer
	from pathlib import Path
	import json
	import time
	import sys

	# Add paths
	sys.path.append(str(Path(__file__).parent / 'training' / 'encoder'))
	sys.path.append(str(Path(__file__).parent / 'training' / 'distillation'))

	from model import create_encoder
	from models import ClassificationExpert

	class HybridMoESystem:
	"""Hierarchical Hybrid MoE: OLMoE Mother + Micro-Experts"""

	def __init__(
	self,
	olmoe_model_path: str,
	encoder_path: Path,
	experts_dir: Path,
	device: str = "auto"
	):
	self.device = device

	# Load OLMoE Mother
	print("🔄 Loading OLMoE-0125 Mother...")
	self.olmoe = AutoModelForCausalLM.from_pretrained(
	olmoe_model_path,
	trust_remote_code=True,
	torch_dtype=torch.float16,
	device_map=device
	)
	self.tokenizer = AutoTokenizer.from_pretrained(olmoe_model_path)
	print("✓ OLMoE loaded")

	# Load Shared Encoder
	print("\n🔄 Loading Shared Encoder...")
	self.encoder_tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
	self.encoder = create_encoder({
	'vocab_size': self.encoder_tokenizer.vocab_size,
	'embed_dim': 256,
	'num_layers': 4,
	'num_heads': 8,
	'ff_dim': 1024,
	'dropout': 0.1,
	'max_seq_len': 128
	})
	self.encoder.load_state_dict(torch.load(encoder_path, map_location='cpu'))
	self.encoder.eval()
	print("✓ Encoder loaded")

	# Load Micro-Experts
	print("\n🔄 Loading Micro-Experts...")
	self.experts = self._load_experts(experts_dir)
	print(f"✓ {len(self.experts)} experts loaded")

	def _load_experts(self, experts_dir: Path):
	"""Load all trained micro-experts"""
	experts = {}

	# Grammar Expert
	grammar_path = experts_dir / 'grammar_expert_v2.pt'
	if grammar_path.exists():
	expert = ClassificationExpert(256, 256, 2, 3, 0.1)
	expert.load_state_dict(torch.load(grammar_path, map_location='cpu'))
	expert.eval()
	experts['grammar'] = {
	'model': expert,
	'name': 'Grammar & Syntax',
	'threshold': 0.7,
	'hint_template': "[GRAMMAR EXPERT]: {advice}"
	}

	# Math Expert
	math_path = experts_dir / 'mathematical_reasoning_expert.pt'
	if math_path.exists():
	expert = ClassificationExpert(256, 256, 1, 3, 0.1)
	expert.load_state_dict(torch.load(math_path, map_location='cpu'))
	expert.eval()
	experts['math'] = {
	'model': expert,
	'name': 'Mathematical Reasoning',
	'threshold': 0.8,
	'hint_template': "[MATH EXPERT]: {advice}"
	}

	# Sentiment Expert
	sentiment_path = experts_dir / 'sentiment_analysis_expert.pt'
	if sentiment_path.exists():
	expert = ClassificationExpert(256, 256, 4, 3, 0.1)
	expert.load_state_dict(torch.load(sentiment_path, map_location='cpu'))
	expert.eval()
	experts['sentiment'] = {
	'model': expert,
	'name': 'Sentiment Analysis',
	'threshold': 0.7,
	'hint_template': "[SENTIMENT EXPERT]: {advice}"
	}

	return experts

	def encode_prompt(self, prompt: str):
	"""Encode prompt with shared encoder"""
	inputs = self.encoder_tokenizer(
	prompt,
	padding='max_length',
	truncation=True,
	max_length=128,
	return_tensors='pt'
	)

	with torch.no_grad():
	embedding = self.encoder(inputs['input_ids'], inputs['attention_mask'])

	return embedding

	def run_experts(self, embedding):
	"""Run all experts in parallel"""
	results = []

	for expert_id, expert_data in self.experts.items():
	model = expert_data['model']

	with torch.no_grad():
	output_emb, confidence, logits = model(embedding)

	# Get confidence score
	probs = torch.softmax(logits, dim=-1)
	max_conf = probs.max().item()
	predicted_class = probs.argmax().item()

	results.append({
	'id': expert_id,
	'name': expert_data['name'],
	'confidence': max_conf,
	'prediction': predicted_class,
	'threshold': expert_data['threshold'],
	'template': expert_data['hint_template']
	})

	return results

	def generate_hints(self, expert_results, prompt):
	"""Generate expert hints for high-confidence activations"""
	hints = []

	for result in expert_results:
	if result['confidence'] > result['threshold']:
	# Generate domain-specific advice
	if result['id'] == 'grammar':
	advice = "Ensure perfect grammar and sentence structure."
	elif result['id'] == 'math':
	advice = "Provide step-by-step mathematical reasoning."
	elif result['id'] == 'sentiment':
	advice = "Consider emotional tone and sentiment."
	else:
	advice = f"Apply {result['name']} expertise."

	hint = result['template'].format(advice=advice)
	hints.append(hint)

	return hints

	def hybrid_generate(
	self,
	prompt: str,
	max_new_tokens: int = 200,
	temperature: float = 0.7,
	top_k: int = 4
	):
	"""Hybrid generation with expert hint injection"""

	print(f"\n{'='*70}")
	print(f"🎯 Hybrid Inference")
	print(f"{'='*70}")
	print(f"Prompt: {prompt}\n")

	# Step 1: Encode prompt
	print("1️⃣ Encoding prompt...")
	embedding = self.encode_prompt(prompt)

	# Step 2: Run experts
	print("2️⃣ Running micro-experts...")
	expert_results = self.run_experts(embedding)

	# Step 3: Filter by confidence
	active_experts = [
	e for e in expert_results
	if e['confidence'] > e['threshold']
	]

	print(f" Active experts: {len(active_experts)}/{len(expert_results)}")
	for exp in active_experts:
	print(f" - {exp['name']}: {exp['confidence']:.2%}")

	# Step 4: Select top-k
	top_experts = sorted(
	active_experts,
	key=lambda x: x['confidence'],
	reverse=True
	)[:top_k]

	# Step 5: Generate hints
	print("\n3️⃣ Generating expert hints...")
	hints = self.generate_hints(top_experts, prompt)

	# Step 6: Inject into system prompt
	if hints:
	hint_text = "\n".join(hints)
	enriched_prompt = f"""{hint_text}

	User: {prompt}
	Assistant:"""
	else:
	enriched_prompt = f"User: {prompt}\nAssistant:"

	print(f"\n4️⃣ OLMoE generation...")

	# Step 7: Generate with OLMoE
	inputs = self.tokenizer(enriched_prompt, return_tensors="pt").to(self.olmoe.device)

	start = time.time()
	outputs = self.olmoe.generate(
	**inputs,
	max_new_tokens=max_new_tokens,
	temperature=temperature,
	do_sample=True
	)
	elapsed = time.time() - start

	response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

	# Extract assistant response
	if "Assistant:" in response:
	response = response.split("Assistant:")[-1].strip()

	tokens_generated = len(outputs[0]) - len(inputs['input_ids'][0])
	tokens_per_sec = tokens_generated / elapsed

	print(f"\n{'='*70}")
	print(f"📊 Results")
	print(f"{'='*70}")
	print(f"Response: {response}\n")
	print(f"⚡ Performance:")
	print(f" Tokens: {tokens_generated}")
	print(f" Time: {elapsed:.2f}s")
	print(f" Speed: {tokens_per_sec:.1f} t/s")
	print(f" Active Experts: {len(top_experts)}")
	print(f"{'='*70}\n")

	return {
	'response': response,
	'active_experts': top_experts,
	'tokens_per_sec': tokens_per_sec,
	'hints': hints
	}


	def main():
	print("="*70)
	print("🚀 Hybrid MoE Inference System")
	print(" OLMoE-0125 + Micro-Experts")
	print("="*70)

	# Initialize system
	system = HybridMoESystem(
	olmoe_model_path="models/olmoe-0125",
	encoder_path=Path("models/encoder_bootstrap.pt"),
	experts_dir=Path("models"),
	device="auto"
	)

	# Test prompts
	test_prompts = [
	"What is 15 + 27?",
	"Fix this sentence: Me and him goes to the store yesterday",
	"Write a Python function to calculate factorial",
	"This movie was absolutely terrible and I hated it",
	]

	for prompt in test_prompts:
	result = system.hybrid_generate(prompt)
	time.sleep(1) # Brief pause between tests


	if __name__ == "__main__":
	main()