Download code/integration_example.py from SetMD/beastbullet-experts: direct link, hf CLI and curl.
- Browser
- Download file 15.6 kB
-
https://huggingface.co/SetMD/beastbullet-experts/resolve/main/code/integration_example.py
- Command line
-
hf download hf://SetMD/beastbullet-experts/code/integration_example.py
-
curl -L -o integration_example.py https://huggingface.co/SetMD/beastbullet-experts/resolve/main/code/integration_example.py
15.6 kB
| #!/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() | |