"""Run upstream torch codec (CPU) on texts + random ids; dump ids/codes/zhats/decoded to w/ref/<task>.npz."""
import sys, os, json, numpy as np, torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from paths import WORK
from ref_torch import load_ref
from texts import *
from transformers import AutoTokenizer
task = sys.argv[1]; d = f"{WORK}/up/{task}"
torch.set_grad_enabled(False); torch.manual_seed(0)
m, cfg = load_ref(d)
tok = AutoTokenizer.from_pretrained(f"{WORK}/tok/{'gpt2' if task=='xsum' else 'Qwen3-0.6B'}")
pt = cfg.tokenizer.pre_tokenizer
EOS = tok.convert_tokens_to_ids(pt.special_tokens.eos_token)
PAD = cfg.model.prefix_pad_token_id
L = cfg.tokenizer.vqvae.context_length
texts = (NEWS if task == "xsum" else PLAIN)
ids = [tok.encode(t)[: L - 1] + [EOS] for t in texts]
ids = [i + [PAD] * (L - len(i)) for i in ids]
rnd = torch.randint(0, 50000, (4, L)).tolist()
x = torch.tensor(ids + rnd)
z = m.encode_z(x)
idx, zh = m.vq.z_to_idx(z, return_zhats=True)
dec = [m.decoder(h).argmax(-1) for h in zh]
out = dict(ids=x.numpy(), z=z.numpy(), idx=idx.numpy(), zh=torch.stack(zh).numpy(), dec=torch.stack(dec).numpy(),
           dec_full=m.decode(idx).argmax(-1).numpy())
np.savez(f"{WORK}/ref_{task}.npz", **out)
print(task, "idx", idx.shape, "EOS", EOS, "PAD", PAD, "L", L)
print("torch roundtrip exact tokens (text rows):", (out["dec"][-1][:len(texts)] == out["ids"][:len(texts)]).mean())