"""OLMoE-1B-7B with streamed 4-bit experts: decode on 4 different prompts (cold cache first), decode-path and prefill
perplexity. Usage: olmoe_bench.py TAG CACHE_GIB [--no-mps-commit]"""
import json, os, sys, time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tag, cache_gib = sys.argv[1], float(sys.argv[2])
ck = "/private/tmp/ssd-moe-local/olmoe-packed"
opts = {"cache_bytes": int(cache_gib * 2**30), "quantized_experts": "/private/tmp/ssd-moe-local/olmoe-experts-q4", "io_threads": 8}
if "--no-mps-commit" in sys.argv: opts["mps_commit"] = False
tok = AutoTokenizer.from_pretrained(ck)
model = AutoModelForCausalLM.from_pretrained(ck, dtype=torch.bfloat16, device_map="mps", experts_offload=opts)
PROMPTS = [
    "Write a short story about a lighthouse keeper who finds a message in a bottle.",
    "Write a Python function that parses an ISO 8601 date string without using datetime, with tests.",
    "Explain the causes of the French Revolution to a high school student.",
    "Translate into French and explain the grammar: 'The children would have gone if it had not rained.'",
]
def chat(p):
    return tok.apply_chat_template([{"role": "user", "content": p}], add_generation_prompt=True, return_tensors="pt", return_dict=True)["input_ids"].to("mps")
runs = []
for p in PROMPTS:
    ids = chat(p); times = []
    class S:
        def put(self, v): times.append(time.perf_counter())
        def end(self): pass
    t0 = time.perf_counter()
    with torch.inference_mode():
        out = model.generate(ids, max_new_tokens=200, min_new_tokens=200, do_sample=False, streamer=S())
    end = time.perf_counter(); first = times[1]
    st = model.ssd_store.stats()
    runs.append({"prompt": p[:40], "decode_tok_s": 199 / (end - first), "text": tok.decode(out[0, ids.shape[1]:])})
text = open("/private/tmp/transformers-moe-ssd-streaming/docs/source/en/philosophy.md").read()
ids = tok(text, return_tensors="pt").input_ids[:, :512].to("mps")
with torch.inference_mode():
    logits = model(ids).logits[0].float()
    prefill_ppl = torch.nn.functional.cross_entropy(logits[:-1], ids[0, 1:]).exp().item()
    out = model(ids[:, :1], use_cache=True); cache = out.past_key_values; nll = []
    for i in range(1, 512):
        nll.append(torch.nn.functional.cross_entropy(out.logits[0, -1:].float(), ids[0, i:i + 1]))
        out = model(ids[:, i:i + 1], past_key_values=cache, use_cache=True)
    decode_ppl = torch.stack(nll).mean().exp().item()
rates = sorted(r["decode_tok_s"] for r in runs)
res = {"tag": tag, "cache_gib": cache_gib, "decode": [round(r["decode_tok_s"], 2) for r in runs], "median": (rates[1] + rates[2]) / 2,
       "prefill_ppl": prefill_ppl, "decode_ppl": decode_ppl, "store": model.ssd_store.stats(), "texts": [r["text"] for r in runs]}
json.dump(res, open(os.path.join(os.path.dirname(os.path.abspath(__file__)), f"res-{tag}.json"), "w"), indent=1)
print(tag, "decode", res["decode"], "median", round(res["median"], 2), "prefill_ppl", round(prefill_ppl, 4), "decode_ppl", round(decode_ppl, 4))