Blame for code/integration_example.py

#!/usr/bin/env python3
"""
BeastBullet Complete Integration Example
Shows how to use experts with TinyLlama and OLMoE for full reasoning pipeline
"""

import torch
import json
from typing import Dict, List, Optional
from pathlib import Path

# ============================================================================
# STEP 1: Load BeastBullet Experts
# ============================================================================

class ExpertLoader:
    """Load and manage BeastBullet expert models."""
    
    def __init__(self, experts_dir: str = "experts/"):
        self.experts_dir = Path(experts_dir)
        self.experts = {}
        self.encoder = None
    
    def load_encoder(self):
        """Load shared encoder for all experts."""
        encoder_path = self.experts_dir / "shared_encoder_v1.0.pt"
        self.encoder = torch.load(encoder_path, map_location="cpu")
        self.encoder.eval()
        print(f"✅ Loaded shared encoder from {encoder_path}")
        return self.encoder
    
    def load_expert(self, expert_name: str):
        """Load a specific expert model."""
        expert_path = self.experts_dir / f"{expert_name}_v1.0.pt"
        if not expert_path.exists():
            raise FileNotFoundError(f"Expert not found: {expert_path}")
        
        expert_model = torch.load(expert_path, map_location="cpu")
        expert_model.eval()
        self.experts[expert_name] = expert_model
        print(f"✅ Loaded {expert_name} from {expert_path}")
        return expert_model
    
    def load_all_experts(self):
        """Load all available expert models."""
        expert_files = list(self.experts_dir.glob("*_expert_v1.0.pt"))
        for expert_file in expert_files:
            expert_name = expert_file.stem.replace("_v1.0", "")
            self.load_expert(expert_name)
        print(f"✅ Loaded {len(self.experts)} experts")
        return self.experts


# ============================================================================
# STEP 2: Integrate with TinyLlama
# ============================================================================

class TinyLlamaIntegration:
    """
    Integrate BeastBullet experts with TinyLlama for synthesis.
    
    Architecture:
    1. Experts analyze query → produce structured evidence
    2. TinyLlama synthesizes evidence → natural language answer
    """
    
    def __init__(self, experts: Dict, encoder, tinyllama_model=None):
        self.experts = experts
        self.encoder = encoder
        self.tinyllama = tinyllama_model or self._load_tinyllama()
    
    def _load_tinyllama(self):
        """Load TinyLlama model from Ollama or Hugging Face."""
        try:
            # Option 1: Use Ollama (recommended)
            import ollama
            print("✅ Using TinyLlama via Ollama")
            return {"type": "ollama", "model": "tinyllama"}
        except ImportError:
            # Option 2: Use Hugging Face
            from transformers import AutoModelForCausalLM, AutoTokenizer
            model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForCausalLM.from_pretrained(model_name)
            print(f"✅ Loaded TinyLlama from Hugging Face: {model_name}")
            return {"type": "hf", "model": model, "tokenizer": tokenizer}
    
    def query(self, user_query: str, config: Optional[Dict] = None) -> Dict:
        """
        Process query through expert → TinyLlama pipeline.
        
        Args:
            user_query: User's question
            config: Optional configuration (from beastbullet_config.py)
        
        Returns:
            Dictionary with answer, confidence, reasoning trace
        """
        config = config or {}
        
        # Step 1: Route to appropriate experts
        selected_experts = self._route_experts(user_query)
        print(f"🔀 Selected experts: {', '.join(selected_experts)}")
        
        # Step 2: Execute experts in parallel
        expert_results = self._execute_experts(user_query, selected_experts)
        
        # Step 3: Validate results
        validated_results = self._validate_results(expert_results)
        
        # Step 4: Synthesize with TinyLlama
        final_answer = self._synthesize_with_tinyllama(
            user_query, 
            validated_results,
            config
        )
        
        return final_answer
    
    def _route_experts(self, query: str) -> List[str]:
        """Simple routing based on query keywords."""
        query_lower = query.lower()
        selected = []
        
        # Math queries
        if any(word in query_lower for word in ["calculate", "math", "%", "percent", "+"]):
            selected.append("math_expert")
        
        # Logic queries
        if any(word in query_lower for word in ["if", "then", "therefore", "logic"]):
            selected.append("logic_expert")
        
        # Code queries
        if any(word in query_lower for word in ["code", "function", "python", "javascript"]):
            selected.append("code_generation_expert")
        
        # Question answering (default)
        if not selected:
            selected.append("question_answering_expert")
        
        # Always add validator
        if "validator_expert" in self.experts:
            selected.append("validator_expert")
        
        return selected
    
    def _execute_experts(self, query: str, expert_names: List[str]) -> Dict:
        """Execute selected experts on query."""
        results = {}
        
        # Encode query
        with torch.no_grad():
            # Simple encoding (you'd use proper tokenization in production)
            query_embedding = torch.randn(1, 256)  # Placeholder
        
        # Execute each expert
        for expert_name in expert_names:
            if expert_name not in self.experts:
                continue
            
            expert = self.experts[expert_name]
            with torch.no_grad():
                output = expert(query_embedding)
                
                # Get prediction and confidence
                probs = torch.softmax(output, dim=-1)
                confidence, prediction = probs.max(dim=-1)
                
                results[expert_name] = {
                    "prediction": prediction.item(),
                    "confidence": confidence.item(),
                    "output": output.tolist()
                }
        
        return results
    
    def _validate_results(self, expert_results: Dict) -> Dict:
        """Validate expert results using validator."""
        if "validator_expert" not in expert_results:
            return expert_results
        
        validator_result = expert_results["validator_expert"]
        
        # If validator says INVALID, flag it
        if validator_result["prediction"] == 1:  # INVALID
            expert_results["validation_warning"] = "Validator detected potential issues"
        
        return expert_results
    
    def _synthesize_with_tinyllama(
        self, 
        query: str, 
        expert_results: Dict,
        config: Dict
    ) -> Dict:
        """Synthesize expert results into natural language using TinyLlama."""
        
        # Build prompt with expert evidence
        prompt = self._build_synthesis_prompt(query, expert_results)
        
        # Generate with TinyLlama
        if self.tinyllama["type"] == "ollama":
            answer = self._generate_ollama(prompt, config)
        else:
            answer = self._generate_hf(prompt, config)
        
        # Compile final result
        return {
            "query": query,
            "answer": answer,
            "expert_results": expert_results,
            "experts_used": list(expert_results.keys()),
            "confidence": self._calculate_overall_confidence(expert_results),
            "reasoning_trace": self._extract_reasoning_trace(expert_results)
        }
    
    def _build_synthesis_prompt(self, query: str, expert_results: Dict) -> str:
        """Build prompt for TinyLlama with expert evidence."""
        prompt = f"""You are a helpful AI assistant. Answer the following question based on expert analysis.

Question: {query}

Expert Analysis:
"""
        
        for expert_name, result in expert_results.items():
            if expert_name == "validator_expert":
                continue
            prompt += f"\n- {expert_name}: Confidence {result['confidence']:.2%}"
        
        if "validation_warning" in expert_results:
            prompt += f"\n\nWarning: {expert_results['validation_warning']}"
        
        prompt += "\n\nProvide a clear, accurate answer with reasoning:"
        
        return prompt
    
    def _generate_ollama(self, prompt: str, config: Dict) -> str:
        """Generate response using Ollama."""
        import ollama
        
        response = ollama.generate(
            model=self.tinyllama["model"],
            prompt=prompt,
            options={
                "temperature": config.get("temperature", 0.3),
                "top_p": config.get("top_p", 0.9),
                "num_predict": config.get("max_tokens", 512),
            }
        )
        
        return response["response"]
    
    def _generate_hf(self, prompt: str, config: Dict) -> str:
        """Generate response using Hugging Face model."""
        tokenizer = self.tinyllama["tokenizer"]
        model = self.tinyllama["model"]
        
        inputs = tokenizer(prompt, return_tensors="pt")
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=config.get("max_tokens", 512),
                temperature=config.get("temperature", 0.3),
                top_p=config.get("top_p", 0.9),
                do_sample=True,
            )
        
        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # Extract answer after prompt
        answer = answer[len(prompt):].strip()
        
        return answer
    
    def _calculate_overall_confidence(self, expert_results: Dict) -> float:
        """Calculate overall confidence from expert results."""
        confidences = [r["confidence"] for r in expert_results.values()]
        return sum(confidences) / len(confidences) if confidences else 0.0
    
    def _extract_reasoning_trace(self, expert_results: Dict) -> List[str]:
        """Extract reasoning trace from expert results."""
        trace = []
        for expert_name, result in expert_results.items():
            trace.append(f"{expert_name}: {result['confidence']:.2%} confidence")
        return trace


# ============================================================================
# STEP 3: Integrate with OLMoE (Advanced)
# ============================================================================

class OLMoEIntegration:
    """
    Integrate BeastBullet experts with OLMoE for hybrid MoE system.
    
    Architecture:
    1. BeastBullet experts handle specialized tasks (math, logic, etc.)
    2. OLMoE handles general language understanding and synthesis
    3. ISL router decides which system to use
    """
    
    def __init__(self, experts: Dict, encoder, olmoe_model=None):
        self.experts = experts
        self.encoder = encoder
        self.olmoe = olmoe_model or self._load_olmoe()
    
    def _load_olmoe(self):
        """Load OLMoE model."""
        try:
            # Option 1: Use Ollama
            import ollama
            print("✅ Using OLMoE via Ollama")
            return {"type": "ollama", "model": "olmoe"}
        except:
            # Option 2: Use Hugging Face
            from transformers import AutoModelForCausalLM, AutoTokenizer
            model_name = "allenai/OLMoE-1B-7B-0125-Instruct"
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForCausalLM.from_pretrained(model_name)
            print(f"✅ Loaded OLMoE from Hugging Face: {model_name}")
            return {"type": "hf", "model": model, "tokenizer": tokenizer}
    
    def query(self, user_query: str, config: Optional[Dict] = None) -> Dict:
        """
        Hybrid query processing:
        - Use BeastBullet experts for specialized tasks
        - Use OLMoE for general language tasks
        """
        config = config or {}
        
        # Determine if query needs specialized experts
        needs_experts = self._needs_specialized_experts(user_query)
        
        if needs_experts:
            # Use BeastBullet experts
            print("🧠 Using BeastBullet experts for specialized reasoning")
            return self._query_with_experts(user_query, config)
        else:
            # Use OLMoE directly
            print("🤖 Using OLMoE for general language task")
            return self._query_with_olmoe(user_query, config)
    
    def _needs_specialized_experts(self, query: str) -> bool:
        """Determine if query needs specialized expert processing."""
        specialized_keywords = [
            "calculate", "math", "logic", "code", "validate",
            "analyze", "reason", "prove", "deduce"
        ]
        return any(kw in query.lower() for kw in specialized_keywords)
    
    def _query_with_experts(self, query: str, config: Dict) -> Dict:
        """Process query using BeastBullet experts + OLMoE synthesis."""
        # Similar to TinyLlama integration but using OLMoE for synthesis
        # (Implementation similar to TinyLlamaIntegration.query)
        pass
    
    def _query_with_olmoe(self, query: str, config: Dict) -> Dict:
        """Process query using OLMoE directly."""
        # Direct OLMoE generation for general tasks
        pass


# ============================================================================
# STEP 4: Complete Usage Example
# ============================================================================

def main():
    """Complete example of using BeastBullet with TinyLlama."""
    
    print("=" * 70)
    print("BeastBullet + TinyLlama Integration Example")
    print("=" * 70)
    
    # Step 1: Load experts
    print("\n📦 Loading BeastBullet experts...")
    loader = ExpertLoader("experts/")
    loader.load_encoder()
    loader.load_all_experts()
    
    # Step 2: Create integration
    print("\n🔗 Creating TinyLlama integration...")
    integration = TinyLlamaIntegration(
        experts=loader.experts,
        encoder=loader.encoder
    )
    
    # Step 3: Load configuration
    print("\n⚙️  Loading Sonnet-quality configuration...")
    with open("configs/config_sonnet_quality.json", "r") as f:
        config = json.load(f)
    
    # Step 4: Query examples
    print("\n" + "=" * 70)
    print("Example Queries")
    print("=" * 70)
    
    queries = [
        "What is 15% of 240?",
        "If all cats are mammals, and Fluffy is a cat, what is Fluffy?",
        "Write a Python function to calculate factorial",
    ]
    
    for i, query in enumerate(queries, 1):
        print(f"\n{'='*70}")
        print(f"Query {i}: {query}")
        print("=" * 70)
        
        result = integration.query(query, config.get("synthesis", {}))
        
        print(f"\n📝 Answer: {result['answer']}")
        print(f"🎯 Confidence: {result['confidence']:.2%}")
        print(f"🧠 Experts used: {', '.join(result['experts_used'])}")
        print(f"\n📊 Reasoning trace:")
        for step in result['reasoning_trace']:
            print(f"  - {step}")
    
    print("\n" + "=" * 70)
    print("✅ Integration complete!")
    print("=" * 70)


if __name__ == "__main__":
    main()