| |
| """ |
| TinyLlama Wrapper for LEGO-MoE |
| Simple interface to TinyLlama via Ollama |
| """ |
|
|
| import subprocess |
| import json |
| from typing import Optional |
|
|
| class TinyLlamaWrapper: |
| """Wrapper for TinyLlama-1.1B via Ollama""" |
| |
| def __init__(self, model_name: str = "tinyllama"): |
| self.model_name = model_name |
| self._verify_model() |
| |
| def _verify_model(self): |
| """Check if TinyLlama is available""" |
| try: |
| result = subprocess.run( |
| ["ollama", "list"], |
| capture_output=True, |
| text=True, |
| timeout=5 |
| ) |
| |
| if self.model_name not in result.stdout: |
| print(f"⚠️ {self.model_name} not found. Pulling...") |
| self._pull_model() |
| else: |
| print(f"✓ {self.model_name} available") |
| |
| except Exception as e: |
| print(f"Error checking model: {e}") |
| |
| def _pull_model(self): |
| """Download TinyLlama if not present""" |
| try: |
| print(f"Downloading {self.model_name}...") |
| subprocess.run( |
| ["ollama", "pull", self.model_name], |
| check=True |
| ) |
| print(f"✓ {self.model_name} downloaded") |
| except Exception as e: |
| print(f"Error downloading model: {e}") |
| raise |
| |
| def generate( |
| self, |
| prompt: str, |
| temperature: float = 0.7, |
| max_tokens: int = 256, |
| stop: Optional[list] = None |
| ) -> str: |
| """ |
| Generate response from TinyLlama |
| |
| Args: |
| prompt: Input prompt |
| temperature: Sampling temperature (0.0-1.0) |
| max_tokens: Maximum tokens to generate |
| stop: Stop sequences |
| |
| Returns: |
| Generated text |
| """ |
| try: |
| |
| cmd = [ |
| "ollama", "run", self.model_name, |
| prompt |
| ] |
| |
| |
| result = subprocess.run( |
| cmd, |
| capture_output=True, |
| text=True, |
| timeout=15 |
| ) |
| |
| if result.returncode == 0: |
| return result.stdout.strip() |
| else: |
| return f"Error: {result.stderr}" |
| |
| except subprocess.TimeoutExpired: |
| return "Error: Generation timeout" |
| except Exception as e: |
| return f"Error: {str(e)}" |
| |
| def generate_with_context( |
| self, |
| query: str, |
| expert_outputs: dict, |
| temperature: float = 0.7 |
| ) -> str: |
| """ |
| Generate response with expert context |
| |
| Args: |
| query: User query |
| expert_outputs: Dict of expert name -> output |
| temperature: Sampling temperature |
| |
| Returns: |
| Generated response |
| """ |
| |
| context = f"Query: {query}\n\n" |
| |
| if expert_outputs: |
| context += "Expert Analysis:\n" |
| for expert_name, output in expert_outputs.items(): |
| context += f"- {expert_name}: {output}\n" |
| context += "\n" |
| |
| context += "Provide a clear, concise response:\n" |
| |
| |
| return self.generate(context, temperature=temperature) |
|
|
|
|
| def test_tinyllama(): |
| """Test TinyLlama wrapper""" |
| print("="*70) |
| print("Testing TinyLlama Wrapper") |
| print("="*70) |
| |
| llama = TinyLlamaWrapper() |
| |
| |
| print("\nTest 1: Simple query") |
| response = llama.generate("What is 2+2?", temperature=0.1, max_tokens=50) |
| print(f"Response: {response}") |
| |
| |
| print("\nTest 2: With expert context") |
| expert_outputs = { |
| "Math Expert": "4" |
| } |
| response = llama.generate_with_context( |
| "What is 2+2?", |
| expert_outputs, |
| temperature=0.1 |
| ) |
| print(f"Response: {response}") |
| |
| print("\n" + "="*70) |
| print("✓ TinyLlama wrapper working!") |
| print("="*70) |
|
|
|
|
| if __name__ == "__main__": |
| test_tinyllama() |
|
|
