# Runs inside a Kaggle kernel (CPU ok). Inputs: /kaggle/input/<SRC>/generator_chat.pt ; outputs ncm_fp16w.onnx + report.
import os, sys, json, types, time
os.environ["NCM_ROOT"]="/kaggle/working/w/ncm"
import torch, torch.nn as nn, numpy as np, ncmchat as N
import torch.nn.functional as _F
_orig_interp=_F.interpolate
def _area_mat(L,o):
    m=np.zeros((o,L),dtype=np.float32)
    for i in range(o):
        a=(i*L)//o; b=-((-(i+1)*L)//o); m[i,a:b]=1.0/(b-a)
    return torch.from_numpy(m)
def _lm(L,o):
    m=np.zeros((o,L),dtype=np.float32)
    for i in range(o):
        s=max((i+0.5)*L/o-0.5,0.0); i0=min(int(s),L-1); i1=min(i0+1,L-1); w=s-i0
        m[i,i0]+=1-w; m[i,i1]+=w
    return torch.from_numpy(m)
def _interp(x,size=None,scale_factor=None,mode="nearest",**kw):
    if mode=="area" and x.dim()==3 and size is not None:
        o=size[0] if isinstance(size,(tuple,list)) else size
        return torch.matmul(x,_area_mat(x.shape[-1],o).t())
    if mode=="linear" and x.dim()==3 and size is not None and not kw.get("align_corners"):
        o=size[0] if isinstance(size,(tuple,list)) else size
        return torch.matmul(x,_lm(x.shape[-1],o).t())
    return _orig_interp(x,size=size,scale_factor=scale_factor,mode=mode,**kw)
_chk=torch.randn(2,5,48); assert torch.allclose(_orig_interp(_chk,size=(26,),mode="area"),_interp(_chk,size=(26,),mode="area"),atol=1e-5)
_F.interpolate=_interp
SRC=os.environ.get("SRC","ncm-chat-smoke")
cands=[f"/kaggle/input/{SRC}/generator_chat.pt"]+[os.path.join(r,"generator_chat.pt") for r,_,f in os.walk("/kaggle/input") if "generator_chat.pt" in f]
gp=[c for c in cands if os.path.exists(c)][0]; print("using",gp,flush=True)
PEFF=int(os.environ.get("PEFF","128"))
import re
_c=open("ck/config.yaml").read(); _c=re.sub(r"prefix_len: 512","prefix_len: %d"%PEFF,_c); open("ck/config_p.yaml","w").write(_c)
CFG,tk,net=N.load("ck",cfg_path="ck/config_p.yaml",dev="cpu",load_gen=False); net.eval()
sd=torch.load(gp,map_location="cpu",weights_only=False)["model"]
own=net.state_dict(); sdf={k:v.float() for k,v in sd.items() if k in own and own[k].shape==v.shape}
print("skipped keys:",[k for k in sd if k not in sdf],flush=True)
net.load_state_dict(sdf,strict=False); assert all(k in sdf for k in own if "attn_mask" not in k and not k.startswith("rope")), "missing weights"
codec=N.Codec(CFG,tk,net); offs=[0]
for l in net.levels: offs.append(offs[-1]+l)
class G(nn.Module):
    def __init__(s):
        super().__init__(); s.net=net; s.vq=tk.vqvae.vq; s.dec=tk.vqvae.decoder; s.enc=tk.vqvae.encoder
    def forward(s, prefix, noise, temp):
        st={"i":0}
        def samp(self_, logits, temperature, top_k, top_p):
            i=st["i"]; st["i"]+=1
            nz=noise[:, offs[i]:offs[i+1], :]
            v,_=torch.topk(logits,20,dim=-1)
            lg=torch.where(logits<v[...,-1:],torch.full_like(logits,-1e9),logits)
            return (lg+temp*nz).argmax(-1)
        s.net._sample_level=types.MethodType(samp,s.net)
        g=s.net.generate(None,prefix,temperature=1.0,top_k=20,top_p=0.9)
        _,zl=s.vq.idx_to_zhats(g,return_zlevels=True)
        return s.dec(zl[-1]).argmax(-1)
m=G().eval()
tests=["Hi, how are you?","What is the capital of France?","Tell me a joke.","What do you like to do for fun?","Can you recommend a good book?","How do I boil an egg?","I'm feeling sad today.","What is the meaning of life?"]
prompts=[N.fmt_prompt([("user",t)]) for t in tests]
pre=torch.tensor([codec.prefix_ids(p) for p in prompts])
Z=torch.zeros(1,offs[-1],4096); T0=torch.tensor(0.0)
def txt(ids):
    ids=list(ids)
    if codec.EOS in ids: ids=ids[:ids.index(codec.EOS)]
    return codec.pre.decode([t for t in ids if t not in (codec.EOS,codec.PAD)],skip_special_tokens=True).strip()
with torch.no_grad(): ref=[m(pre[i:i+1],Z,T0)[0].tolist() for i in range(len(tests))]
print("torch greedy:",[txt(r) for r in ref],flush=True)
os.makedirs("/kaggle/working/out",exist_ok=True)
f32="/kaggle/working/ncm_fp32.onnx"
with torch.no_grad(): torch.onnx.export(m,(pre[:1],Z,T0),f32,input_names=["prefix","noise","temp"],output_names=["ids"],opset_version=17,dynamo=False)
print("fp32 onnx MB",os.path.getsize(f32)/1e6,flush=True)
import onnxruntime as ort
from onnxruntime.quantization import quantize_dynamic, QuantType
import onnx
from onnx import numpy_helper, helper
mo=onnx.load(f32); g=mo.graph; inits={i.name:i for i in g.initializer}; newn=[]; ng=0
for n in g.node:
    if n.op_type=="Gemm" and n.input[1] in inits and (len(n.input)<3 or n.input[2] in inits):
        at={a.name:helper.get_attribute_value(a) for a in n.attribute}
        W=numpy_helper.to_array(inits[n.input[1]]); W=(W.T if at.get("transB",0) else W)*at.get("alpha",1.0)
        wn=n.input[1]+"_mm"; g.initializer.append(numpy_helper.from_array(W.astype("float32"),wn))
        mm=n.output[0]+"_mmout" if len(n.input)>=3 else n.output[0]
        newn.append(helper.make_node("MatMul",[n.input[0],wn],[mm],name=n.name+"_mm"))
        if len(n.input)>=3: newn.append(helper.make_node("Add",[mm,n.input[2]],[n.output[0]],name=n.name+"_add"))
        ng+=1
    else: newn.append(n)
# remove Identity nodes fed by initializers (they hid the weights from the int8 quantizer)
inits_all={i.name for i in g.initializer}
alias={}
final=[]
for n in newn:
    if n.op_type=="Identity" and (n.input[0] in inits_all or n.input[0] in alias):
        alias[n.output[0]]=n.input[0]
    else: final.append(n)
def res(x):
    while x in alias: x=alias[x]
    return x
for n in final:
    for j,x in enumerate(n.input): n.input[j]=res(x)
for o in g.output:
    if o.name in alias: raise SystemExit("identity feeds graph output")
del g.node[:]; g.node.extend(final); print("gemm->matmul",ng,"identity removed",len(alias),flush=True)
f32m="/kaggle/working/ncm_fp32_mm.onnx"; onnx.save(mo,f32m)
i8="/kaggle/working/out/ncm_fp16w.onnx"   # weights stored fp16 + Cast->fp32 at load (compute stays fp32)
mo=onnx.load(f32m); g=mo.graph; casts=[]; nconv=0
for init in list(g.initializer):
    if init.data_type==onnx.TensorProto.FLOAT and int(np.prod(init.dims))>=1024:
        arr=numpy_helper.to_array(init); h=numpy_helper.from_array(arr.astype(np.float16),init.name+"_h")
        casts.append(helper.make_node("Cast",[init.name+"_h"],[init.name],to=onnx.TensorProto.FLOAT,name=init.name+"_cast"))
        g.initializer.remove(init); g.initializer.append(h); nconv+=1
nodes=list(g.node); del g.node[:]; g.node.extend(casts+nodes)
onnx.save(mo,i8); print("fp16-weights onnx MB",os.path.getsize(i8)/1e6,"converted",nconv,flush=True)
res={"files_MB":{"fp32":os.path.getsize(f32)/1e6,"fp16w":os.path.getsize(i8)/1e6},"cases":[]}
for name,path in [("fp32",f32m),("fp16w",i8)]:
    s=ort.InferenceSession(path,providers=["CPUExecutionProvider"]); outs=[]
    t=time.time()
    for i in range(len(tests)):
        o=s.run(None,{"prefix":pre[i:i+1].numpy(),"noise":Z.numpy(),"temp":np.array(0.0,dtype=np.float32)})[0][0].tolist()
        outs.append(o)
    dt=(time.time()-t)/len(tests); res[name]={"sec_per_reply_cpu":dt,"exact_match_vs_torch":sum(o==r for o,r in zip(outs,ref)),"replies":[txt(o) for o in outs]}
    print(name,"exact-match",res[name]["exact_match_vs_torch"],"/",len(tests),"s/reply",round(dt,2),flush=True)
    for t_,o in zip(tests,outs): print("  ",t_,"->",txt(o),flush=True)
res["torch_replies"]=[txt(r) for r in ref]
# tokenizer parity vectors for the browser
res["tok_vectors"]=[{"prompt":p,"ids":codec.enc(p),"prefix_ids":codec.prefix_ids(p)[:len(codec.enc(p))]} for p in prompts[:4]]
res["EOS"]=codec.EOS; res["PAD"]=codec.PAD; res["ppad"]=codec.ppad; res["P"]=codec.P; res["offs"]=offs
json.dump(res,open("/kaggle/working/out/report.json","w"),indent=1)
print("done",flush=True)