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")