#!/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()