Upload moe.py with huggingface_hub

moe.py

@@ -0,0 +1,44 @@
+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")