beastbullet-experts / code /integration_example.py
SetMD's picture
Add TinyLlama and OLMoE integration guide + example code
021efe8
Raw History Blame Contribute Delete
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()