File size: 4,300 Bytes
ae1e90a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | #!/usr/bin/env python3
"""
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:
# Build Ollama command
cmd = [
"ollama", "run", self.model_name,
prompt
]
# Run generation with shorter timeout
result = subprocess.run(
cmd,
capture_output=True,
text=True,
timeout=15 # Reduced from 30 to 15 seconds
)
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
"""
# Build context
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"
# Generate
return self.generate(context, temperature=temperature)
def test_tinyllama():
"""Test TinyLlama wrapper"""
print("="*70)
print("Testing TinyLlama Wrapper")
print("="*70)
llama = TinyLlamaWrapper()
# Test 1: Simple generation
print("\nTest 1: Simple query")
response = llama.generate("What is 2+2?", temperature=0.1, max_tokens=50)
print(f"Response: {response}")
# Test 2: With expert context
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()
|