Upload simple_router.py with huggingface_hub

simple_router.py

@@ -0,0 +1,114 @@
+"""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:]}')