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