File size: 4,046 Bytes
df441e9 | 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 | """Simple Router — prompt classification + expert routing.
Classifies coding/math/chat then loads the right LoRA expert.
"""
import torch
import re
from transformers import AutoModelForCausalLM, AutoTokenizer
# ─── Classification ───
def classify_prompt(text: str) -> str:
"""Classify prompt into domain using keyword matching."""
text_lower = text.lower()
coding_kw = [
'def ', 'function', 'python', 'code', 'bug', 'debug', 'api',
'import', 'class ', 'algorithm', 'implement', 'compile', 'syntax',
'javascript', 'html', 'css', 'sql', 'bash', 'git', 'docker',
'write a', 'program', 'script', 'loop', 'array', 'list', 'dict',
]
math_kw = [
'solve', 'equation', 'derivative', 'integral', 'matrix', 'eigen',
'theorem', 'proof', 'sqrt', 'log', 'sin', 'cos', 'tan', 'sum',
'probability', 'statistic', 'graph', 'vector', 'polynomial',
'x =', 'x=', 'y =', 'calculate', 'compute', 'find the',
]
coding_score = sum(1 for kw in coding_kw if kw in text_lower)
math_score = sum(1 for kw in math_kw if kw in text_lower)
if coding_score > 0 and coding_score >= math_score:
return 'coding'
elif math_score > 0 and math_score > coding_score:
return 'math'
else:
return 'chat'
# ─── Router ───
class ExpertRouter:
def __init__(self, base_model_name: str = 'unsloth/Qwen2.5-1.5B-Instruct'):
self.base_name = base_model_name
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
# Load base model once
print(f'Loading base model: {base_model_name}...')
self.base = AutoModelForCausalLM.from_pretrained(
base_model_name,
torch_dtype=torch.bfloat16,
device_map='auto',
)
self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
# Cache for loaded experts
self._experts = {}
self._current = None
def _load_expert(self, domain: str):
"""Load LoRA adapter for a domain."""
from peft import PeftModel
import copy
if domain in self._experts:
return self._experts[domain]
print(f' Loading expert: {domain}...')
# Load fresh base + adapter each time (merge_and_unload corrupts base)
base_fresh = AutoModelForCausalLM.from_pretrained(
self.base_name,
torch_dtype=torch.bfloat16,
device_map='auto',
)
model = PeftModel.from_pretrained(
base_fresh,
'hotdogs/frankenmoe',
subfolder=domain,
torch_dtype=torch.bfloat16,
)
model = model.merge_and_unload()
self._experts[domain] = model
return model
def generate(self, prompt: str, max_tokens: int = 128) -> tuple:
"""Classify + generate response."""
domain = classify_prompt(prompt)
model = self._load_expert(domain)
inp = self.tokenizer(prompt, return_tensors='pt').to(self.device)
with torch.no_grad():
out = model.generate(
**inp,
max_new_tokens=max_tokens,
do_sample=True,
temperature=0.7,
top_p=0.9,
)
text = self.tokenizer.decode(out[0], skip_special_tokens=True)
return domain, text
# ─── Main ───
if __name__ == '__main__':
router = ExpertRouter()
tests = [
'Write a Python function to reverse a linked list',
'Solve the quadratic equation 2x^2 - 4x + 1 = 0',
'What is the capital of Thailand?',
]
for prompt in tests:
domain, response = router.generate(prompt)
print(f'\n{"="*50}')
print(f'[ROUTE: {domain}]')
print(f'PROMPT: {prompt}')
print(f'OUTPUT: {response[-300:]}')
|