| import torch, joblib |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| from peft import PeftModel |
| from sentence_transformers import SentenceTransformer |
|
|
| BASE = "microsoft/bitnet-b1.58-2B-4T-bf16" |
| ADAPTERS = { |
| "mult": "UlukaDev/bitnet-2digit-mult-expert", |
| "roman": "UlukaDev/bitnet-roman-numeral-expert", |
| } |
| SYS = "You are a careful calculator. Work step by step, then end with exactly 'The answer is X'." |
|
|
| tok = AutoTokenizer.from_pretrained(BASE) |
| if tok.pad_token is None: tok.pad_token = tok.eos_token |
| tok.padding_side = "left" |
|
|
| base = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.bfloat16, device_map="auto") |
| names = list(ADAPTERS) |
| model = PeftModel.from_pretrained(base, ADAPTERS[names[0]], adapter_name=names[0]) |
| for n in names[1:]: |
| model.load_adapter(ADAPTERS[n], adapter_name=n) |
| model.eval() |
|
|
| enc = SentenceTransformer("all-MiniLM-L6-v2") |
| clf = joblib.load("router.joblib") |
|
|
| def route(q): |
| return clf.predict(enc.encode([q], normalize_embeddings=True))[0] |
|
|
| @torch.no_grad() |
| def moe_generate(q, max_new_tokens=250): |
| expert = route(q) |
| model.set_adapter(expert) |
| msgs = [{"role":"system","content":SYS}, {"role":"user","content":q}] |
| enc_in = tok.apply_chat_template(msgs, add_generation_prompt=True, |
| return_tensors="pt", return_dict=True).to(model.device) |
| in_len = enc_in["input_ids"].shape[1] |
| out = model.generate(**enc_in, max_new_tokens=max_new_tokens, do_sample=False) |
| return expert, tok.decode(out[0, in_len:], skip_special_tokens=True) |
|
|
| if __name__ == "__main__": |
| for q in ["What is 34 times 57?", "Convert 1994 to Roman numerals"]: |
| expert, ans = moe_generate(q) |
| print(f"[{expert}] {q}\n -> {ans}\n") |
|
|
