"""Generate text from alpha-er with the PyTorch port.

Pads to block_size and reads the last real position, which the conditional MLP
requires and which is exact: attention is causal, and each token's expert is a
function of its own position.
"""
import sys, os, json, torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from modeling_alpha import AlphaErConfig, AlphaErForCausalLM
from tokenization_alpha import AlphaErTokenizer
from safetensors.torch import load_file

hf_dir = sys.argv[1]
prompts = sys.argv[2:] or ["<|user|>Hello!<|assistant|>"]
temp = float(os.environ.get("TEMP", "0.8"))
topk = int(os.environ.get("TOPK", "40"))
ntok = int(os.environ.get("NTOK", "60"))

cfg_d = json.load(open(f"{hf_dir}/config.json"))
cfg = AlphaErConfig(**{k: v for k, v in cfg_d.items()
                       if k in AlphaErConfig.__init__.__code__.co_varnames})
model = AlphaErForCausalLM(cfg)
model.load_state_dict(load_file(f"{hf_dir}/model.safetensors"), strict=False)
model.eval()
tok = AlphaErTokenizer.from_file(f"{hf_dir}/tokenizer_artifacts.json")
torch.manual_seed(1234)

print(f"alpha-er  step {cfg_d.get('trained_step')}  temp={temp} top_k={topk}\n")
for p in prompts:
    ids = tok.encode(p)
    out = model.generate(torch.tensor([ids]), max_new_tokens=ntok,
                         temperature=temp, top_k=topk)[0].tolist()
    print("=" * 72)
    print("PROMPT:", p)
    print("OUTPUT:", tok.decode(out[len(ids):]).replace("\n", "\\n"))