feat(prefetch): implement prefetcher v2.1 with lookahead-2, hot pinning, adaptive cache, RMSNorm scaling, and routing EMA

c/tools/make_olmoe_real_oracle.py

@@ -0,0 +1,74 @@
+"""Generate reference token IDs for the real OLMoE-1B-7B model.
+
+Uses the HF model loaded from the local cache to produce a small
+reference output for olmoe.exe validation. Saves to ref_olmoe_real.json.
+
+Usage: python tools/make_olmoe_real_oracle.py
+"""
+import json
+import sys
+from pathlib import Path
+
+if sys.platform == "win32":
+    for s in (sys.stdout, sys.stderr):
+        try:
+            s.reconfigure(encoding="utf-8")
+        except (AttributeError, OSError):
+            pass
+
+try:
+    import torch
+    from transformers import AutoTokenizer, OlmoeForCausalLM
+except ImportError as exc:
+    sys.exit(f"Missing deps: {exc}. Run: pip install torch transformers")
+
+MODEL_DIR = (
+    Path(r"C:\Users\egonr\.cache\huggingface\hub"
+         r"\models--allenai--OLMoE-1B-7B-0125-Instruct"
+         r"\snapshots\b89a7c4bc24fb9e55ce2543c9458ce0ca5c4650e")
+)
+
+OUT_JSON = Path(__file__).resolve().parent.parent / "ref_olmoe_real.json"
+
+PROMPT = "The capital of France is"
+MAX_NEW_TOKENS = 12
+
+print(f"Loading tokenizer from {MODEL_DIR} ...")
+tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR))
+
+print("Encoding prompt ...")
+enc = tokenizer(PROMPT, return_tensors="pt")
+prompt_ids = enc["input_ids"][0].tolist()
+print(f"  Prompt IDs ({len(prompt_ids)}): {prompt_ids}")
+
+print(f"Loading OLMoE model from {MODEL_DIR} ...")
+print("  (this will use ~14 GB RAM — please be patient)")
+model = OlmoeForCausalLM.from_pretrained(
+    str(MODEL_DIR),
+    torch_dtype=torch.bfloat16,
+    device_map="cpu",
+    low_cpu_mem_usage=True,
+)
+model.eval()
+print("  Model loaded!")
+
+print(f"Generating {MAX_NEW_TOKENS} tokens ...")
+with torch.no_grad():
+    out = model.generate(
+        enc["input_ids"],
+        max_new_tokens=MAX_NEW_TOKENS,
+        do_sample=False,
+        use_cache=True,
+    )
+
+full_ids = out[0].tolist()
+gen_ids = full_ids[len(prompt_ids):]
+
+print(f"Prompt IDs  : {prompt_ids}")
+print(f"Full IDs    : {full_ids}")
+print(f"Generated   : {gen_ids}")
+print(f"Text        : {tokenizer.decode(gen_ids, skip_special_tokens=True)!r}")
+
+payload = {"prompt_ids": prompt_ids, "full_ids": full_ids}
+OUT_JSON.write_text(json.dumps(payload, indent=2))
+print(f"\nSaved reference to {OUT_JSON}")