Upload folder using huggingface_hub

README.md

@@ -0,0 +1,18 @@
+---
+license: mit
+tags:
+- next-concept-modeling
+- ncm
+- experimental
+---
+
+# NCM chat generator (experimental, rough)
+
+Fine-tuned generator weights (fp16, 177 MB) for a small chat demo built on **Next-Concept Modeling** (Stanford Hazy Research, arXiv 2610.14162, https://github.com/HazyResearch/ncm).
+
+- Base: the released XSum checkpoint at https://huggingface.co/hazyresearch/ncm-checkpoints (frozen multiscale VQ tokenizer plus prefix-conditioned generator). Only the generator was fine-tuned, for 4000 steps on short public dialogue and instruction data (everyday-conversations, dolly, alpaca; 21.6k examples). Held-out code loss 6.31 to 1.47.
+- Load it with the base repo's config, tokenizer and code (`code/ncmchat.py` shows how). `generator_chat.pt` holds a `model` state dict.
+- Prefix window 128 tokens, reply up to 48 tokens, levels [4,6,10,14,20,26,32,40,48].
+- `code/` has the training, ONNX export and parity scripts.
+- Quality is rough: replies are short and often repetitive or wrong. This is a research demo, not an assistant.
+- Browser demo: https://huggingface.co/spaces/danielharkin21/ncm-chat-experimental

code/diag2_body.py

@@ -0,0 +1,110 @@
+# Teacher-forced logit parity: identical prefix + identical codes through torch and ONNX (fp32, fp32 no-opt, int8).
+import os, sys, json, types, re, 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=_F.interpolate
+def _am(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,_am(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(x,size=size,scale_factor=scale_factor,mode=mode,**kw)
+gp=[os.path.join(r,"generator_chat.pt") for r,_,f in os.walk("/kaggle/input") if "generator_chat.pt" in f][0]
+P=128
+c=open("ck/config.yaml").read(); c=re.sub(r"prefix_len: 512","prefix_len: %d"%P,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()
+net.load_state_dict({k:v.float() for k,v in sd.items() if k in own and own[k].shape==v.shape},strict=False)
+codec=N.Codec(CFG,tk,net); offs=[0]
+for l in net.levels: offs.append(offs[-1]+l)
+levels=list(net.levels)
+class GTF(nn.Module):   # generate-path with forced codes; returns logits of every level
+    def __init__(s):
+        super().__init__(); s.net=net
+    def forward(s,prefix,codes):
+        st={"i":0}; L=[]
+        def samp(self_,logits,temperature,top_k,top_p):
+            i=st["i"]; st["i"]+=1; L.append(logits)
+            return codes[:,offs[i]:offs[i+1]]
+        s.net._sample_level=types.MethodType(samp,s.net)
+        s.net.generate(None,prefix,temperature=1.0,top_k=20,top_p=0.9)
+        return torch.cat(L,1)
+m=GTF().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])
+# codes: use the REAL data codes of plausible replies (teacher forcing), via tokenizer encode
+reps=["Hello! How can I help you today?","The capital of France is Paris.","Why did the chicken cross the road? To get to the other side.","I like to read books and go hiking.","I recommend The Hobbit, it is a fun story.","Put it in boiling water for ten minutes.","I am sorry to hear that. I hope you feel better.","Life is about finding what makes you happy."]
+tgt=torch.tensor([codec.target_ids(r) for r in reps])
+with torch.no_grad(): codes=tk.vqvae.encode(tgt)
+# (1) training-path forward logits vs generate-path logits (torch vs torch)
+with torch.no_grad():
+    lg_gen=torch.cat([m(pre[i:i+1],codes[i:i+1]) for i in range(len(tests))],0)
+    x=N.build_inputs(codec,pre,codes,0.0); pm=pre==codec.ppad
+    lg_tr=net(x,prefix_pad_mask=pm)
+d_tr=(lg_gen-lg_tr).abs()
+print("torch generate-path vs training-forward logits: max|d|",float(d_tr.max()),"mean|d|",float(d_tr.mean()),"argmax mismatch",int((lg_gen.argmax(-1)!=lg_tr.argmax(-1)).sum()),"/",lg_gen.shape[0]*lg_gen.shape[1],flush=True)
+os.makedirs("/kaggle/working/out",exist_ok=True)
+_F.interpolate=_interp
+torch.onnx.export(m,(pre[:1],codes[:1]),"/kaggle/working/tf.onnx",input_names=["prefix","codes"],output_names=["logits"],opset_version=17,dynamo=False)
+import onnxruntime as ort, onnx
+from onnx import numpy_helper, helper
+def clean_quant(src,dst):
+    mo=onnx.load(src); g=mo.graph
+    keep=[]
+    for n in g.node:   # Constant nodes -> initializers so the quantizer sees weights
+        if n.op_type=="Constant" and len(n.attribute)==1 and n.attribute[0].name=="value":
+            t=n.attribute[0].t; t.name=n.output[0]; g.initializer.append(t)
+        else: keep.append(n)
+    del g.node[:]; g.node.extend(keep)
+    inits={i.name for i in g.initializer}; alias={}; final=[]
+    for n in g.node:
+        if n.op_type=="Identity" and (n.input[0] in inits or n.input[0] in alias): alias[n.output[0]]=n.input[0]
+        else: final.append(n)
+    def r(x):
+        while x in alias: x=alias[x]
+        return x
+    for n in final:
+        for j,xx in enumerate(n.input): n.input[j]=r(xx)
+    del g.node[:]; g.node.extend(final); onnx.save(mo,"/kaggle/working/tf_clean.onnx")
+    from onnxruntime.quantization import quantize_dynamic, QuantType
+    quantize_dynamic("/kaggle/working/tf_clean.onnx",dst,weight_type=QuantType.QInt8)
+HAVE_I8=True
+try:
+    clean_quant("/kaggle/working/tf.onnx","/kaggle/working/tf_i8.onnx"); print("int8 MB",os.path.getsize("/kaggle/working/tf_i8.onnx")/1e6,flush=True)
+except Exception as e:
+    HAVE_I8=False; print("INT8 QUANT FAILED",repr(e)[:300],flush=True)
+def run(path,opt):
+    so=ort.SessionOptions(); so.graph_optimization_level=opt
+    s=ort.InferenceSession(path,so,providers=["CPUExecutionProvider"])
+    return torch.cat([torch.from_numpy(s.run(None,{"prefix":pre[i:i+1].numpy(),"codes":codes[i:i+1].numpy()})[0]) for i in range(len(tests))],0)
+rep={}
+for name,path,opt in [("onnx_fp32_default","/kaggle/working/tf.onnx",ort.GraphOptimizationLevel.ORT_ENABLE_ALL),("onnx_fp32_noopt","/kaggle/working/tf.onnx",ort.GraphOptimizationLevel.ORT_DISABLE_ALL),("onnx_int8","/kaggle/working/tf_i8.onnx",ort.GraphOptimizationLevel.ORT_ENABLE_ALL)]:
+    if name=="onnx_int8" and not HAVE_I8: continue
+    lg=run(path,opt); d=(lg-lg_gen).abs(); top2=lg_gen.topk(2,-1).values; gap=top2[...,0]-top2[...,1]
+    mism=(lg.argmax(-1)!=lg_gen.argmax(-1))
+    per=[]
+    for k in range(len(levels)):
+        a,b=offs[k],offs[k+1]
+        per.append({"level":k,"max_abs":float(d[:,a:b].max()),"mean_abs":float(d[:,a:b].mean()),"argmax_mismatch":int(mism[:,a:b].sum()),"of":int(mism[:,a:b].numel())})
+    rep[name]={"max_abs":float(d.max()),"mean_abs":float(d.mean()),"logit_std":float(lg_gen.std()),"argmax_mismatch":int(mism.sum()),"of":int(mism.numel()),
+      "gap_at_mismatches_max":float(gap[mism].max()) if mism.any() else None,"gap_at_mismatches_median":float(gap[mism].median()) if mism.any() else None,"per_level":per}
+    print(name,"max|d|",rep[name]["max_abs"],"mean|d|",rep[name]["mean_abs"],"logit std",rep[name]["logit_std"],"argmax mism",rep[name]["argmax_mismatch"],"/",rep[name]["of"],"gap@mism max/med",rep[name]["gap_at_mismatches_max"],rep[name]["gap_at_mismatches_median"],flush=True)
+    for p_ in per: print("   ",p_,flush=True)
+rep["torch_gen_vs_train_forward"]={"max_abs":float(d_tr.max()),"mean_abs":float(d_tr.mean())}
+json.dump(rep,open("/kaggle/working/out/parity.json","w"),indent=1)
+print("done",flush=True)

code/export_onnx_body.py

@@ -0,0 +1,127 @@
+# 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)

code/ncm_chat_train.py

@@ -0,0 +1,102 @@
+import os, subprocess, sys, json, random, time, math
+STEPS=int(os.environ.get("STEPS","300")); BS=int(os.environ.get("BS","24")); LR=float(os.environ.get("LR","1e-4"))
+def sh(c): print("$",c,flush=True); subprocess.run(c,shell=True,check=True)
+os.makedirs("/kaggle/working/w",exist_ok=True); os.chdir("/kaggle/working/w")
+if not os.path.exists("ncm"): sh("git clone -q https://github.com/HazyResearch/ncm.git")
+sh("pip install -q omegaconf hydra-core regex 2>&1 | tail -1")
+os.makedirs("ck",exist_ok=True)
+for f in ["config.yaml","generator.pt","tokenizer.pt"]:
+    if not os.path.exists("ck/"+f): sh(f"curl -sL -o ck/{f} https://huggingface.co/hazyresearch/ncm-checkpoints/resolve/main/xsum/{f}")
+open('ncmchat.py','w').write('"""NCM chat fine-tune: reuse released XSum NCM (frozen multiscale VQ tokenizer + prefix-conditioned generator),\nfine-tune the generator on (conversation prefix -> short reply). Needs HazyResearch/ncm on NCM_ROOT."""\nimport os, sys, random, torch, torch.nn.functional as F\nfrom omegaconf import OmegaConf\nNCM_ROOT = os.environ.get("NCM_ROOT", "/tmp/ncm"); sys.path.insert(0, NCM_ROOT)\nfrom models.next_concept import NextConceptTransformer\nfrom tokenization.vqvae_tokenizer import VQVAETokenizer\n\ndef load(ck, cfg_path=None, dev="cpu", load_gen=True):\n    CFG = OmegaConf.load(cfg_path or os.path.join(ck, "config.yaml"))\n    CFG.tokenizer.vqvae.checkpoint_path = os.path.join(ck, CFG.checkpoints.tokenizer)\n    tk = VQVAETokenizer.from_config(CFG.tokenizer)\n    tk.vqvae = tk.vqvae.to(dev).eval()\n    for p in tk.vqvae.parameters(): p.requires_grad = False\n    net = NextConceptTransformer(vqvae_local=tk.vqvae, config=CFG.model).to(dev)\n    if load_gen:\n        sd = torch.load(os.path.join(ck, CFG.checkpoints.generator), map_location="cpu", weights_only=False)\n        net.load_state_dict(sd["model"], strict=True)\n    return CFG, tk, net\n\ndef fmt_prompt(turns):\n    """turns: list of (role, text) ending with a user turn -> prompt string"""\n    s = "".join(f"{\'User\' if r==\'user\' else \'Assistant\'}: {t}\\n" for r, t in turns)\n    return s + "Assistant:"\n\nclass Codec:\n    def __init__(self, CFG, tk, net):\n        self.tk, self.net, self.P = tk, net, int(CFG.model.prefix_len)\n        self.pre = tk.pre_tokenizer\n        self.EOS, self.PAD = self.pre.eos_token_id, self.pre.pad_token_id\n        self.ppad = int(CFG.model.prefix_pad_token_id)   # pad id used in the generator prefix\n        self.T = tk.vqvae.context_length                  # target length (48)\n        self.hf = self.pre.tokenizer if hasattr(self.pre, "tokenizer") else None\n    def enc(self, text): return self.pre.encode(text)\n    def prefix_ids(self, prompt):\n        ids = self.enc(prompt)[-self.P:]                   # keep most recent context\n        return ids + [self.ppad] * (self.P - len(ids))\n    def target_ids(self, reply):\n        ids = self.enc(" " + reply.strip())[: self.T - 1] + [self.EOS]\n        return ids + [self.EOS] * (self.T - len(ids))\n\n@torch.no_grad()\ndef target_codes(codec, tgt):                              # tgt: LongTensor (B,T) -> idx (B, sum levels)\n    return codec.tk.vqvae.encode(tgt)\n\ndef build_inputs(codec, prefix, idx, corrupt=0.1):\n    net, vq = codec.net, codec.tk.vqvae.vq\n    with torch.no_grad():\n        idc = idx.clone()\n        if corrupt > 0:                                    # robustness: corrupt conditioning codes (as in paper)\n            m = torch.rand_like(idc, dtype=torch.float) < corrupt\n            idc[m] = torch.randint(0, 4096, (int(m.sum()),), device=idc.device)\n        nc = vq.idx_to_ncm_input(idc)                      # cumulative recon fed to levels 1..K-1\n        zp = codec.tk.vqvae.encoder.wte(prefix)[:, :codec.P, :]\n        q0 = net.build_l0_query_input(prefix.shape[0])\n    return torch.cat([zp, q0, nc], dim=1)\n\ndef loss_fn(codec, prefix, tgt, corrupt=0.1):\n    idx = target_codes(codec, tgt)\n    x = build_inputs(codec, prefix, idx, corrupt)\n    pad_mask = prefix == codec.ppad\n    logits = codec.net(x, prefix_pad_mask=None if not hasattr(codec.net, "_prefix_pad_bias") else pad_mask)\n    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]).float(), idx.reshape(-1)), logits, idx\n\n@torch.no_grad()\ndef reply(codec, prompts, temperature=0.8, top_k=20, top_p=0.9, dev="cpu"):\n    net, tk = codec.net, codec.tk\n    pre = torch.tensor([codec.prefix_ids(p) for p in prompts], device=dev)\n    g = net.generate(None, pre, temperature=temperature, top_k=top_k, top_p=top_p)\n    _, zl = tk.vqvae.vq.idx_to_zhats(g, return_zlevels=True)\n    ids = tk.vqvae.decoder(zl[-1]).argmax(-1).tolist()\n    out = []\n    for r in ids:\n        if codec.EOS in r: r = r[: r.index(codec.EOS)]\n        out.append(codec.pre.decode([t for t in r if t not in (codec.EOS, codec.PAD)], skip_special_tokens=True).strip())\n    return out\n')
+sys.path.insert(0,os.getcwd())
+os.environ["NCM_ROOT"]="/kaggle/working/w/ncm"
+import torch, ncmchat as N
+dev="cuda"
+CFG,tk,net=N.load("ck",dev=dev); codec=N.Codec(CFG,tk,net)
+# ---------------- data ----------------
+random.seed(0)
+from datasets import load_dataset
+ex=[]  # (turns list[(role,text)], reply)
+def add_dialog(utts, maxn=6):
+    # utts alternating, first is "user"
+    for i in range(1,len(utts)):
+        if i%2==1:
+            hist=[("user" if j%2==0 else "assistant",utts[j]) for j in range(max(0,i-3),i)]
+            if hist[0][0]=="assistant": hist=hist[1:]
+            ex.append((hist,utts[i]))
+src={}
+def tryload(name,fn):
+    try:
+        n0=len(ex); fn(); src[name]=len(ex)-n0; print("loaded",name,src[name],flush=True)
+    except Exception as e: print("FAILED",name,repr(e)[:200],flush=True)
+def dd():
+    d=load_dataset("roskoN/dailydialog",split="train")
+    for r in d:
+        u=r.get("utterances") or r.get("dialog")
+        add_dialog([x.strip() for x in u])
+def ev():
+    d=load_dataset("HuggingFaceTB/everyday-conversations-llama3.1-2k",split="train_sft")
+    for r in d:
+        m=r["messages"]; add_dialog([x["content"].strip() for x in m if x["role"]!="system"])
+def dolly():
+    d=load_dataset("databricks/databricks-dolly-15k",split="train")
+    for r in d:
+        if not r["context"] and r["response"]: ex.append(([("user",r["instruction"].strip())],r["response"].strip()))
+def alp():
+    d=load_dataset("tatsu-lab/alpaca",split="train")
+    for r in d:
+        if not r["input"]: ex.append(([("user",r["instruction"].strip())],r["output"].strip()))
+tryload("dailydialog",dd); tryload("everyday",ev); tryload("dolly",dolly); tryload("alpaca",alp)
+print("raw examples",len(ex),flush=True)
+def ok(e):
+    h,r=e
+    return 0<len(codec.enc(r))<=40 and "\n" not in r and len(r)>1
+ex=[e for e in ex if ok(e)]
+random.shuffle(ex); print("kept short-reply examples",len(ex),flush=True)
+val=ex[:500]; tr=ex[500:]
+def batch(items):
+    pre=torch.tensor([codec.prefix_ids(N.fmt_prompt(h)) for h,_ in items],device=dev)
+    tgt=torch.tensor([codec.target_ids(r) for _,r in items],device=dev)
+    return pre,tgt
+# ---------------- train ----------------
+opt=torch.optim.AdamW([p for p in net.parameters() if p.requires_grad],lr=LR,weight_decay=0.01,betas=(0.9,0.95))
+scaler=torch.cuda.amp.GradScaler()
+def evalval():
+    net.eval(); tot=0;n=0
+    with torch.no_grad():
+        for i in range(0,len(val),50):
+            pre,tgt=batch(val[i:i+50])
+            with torch.autocast("cuda",dtype=torch.float16): l,_,_=N.loss_fn(codec,pre,tgt,0.0)
+            tot+=l.item()*len(pre); n+=len(pre)
+    net.train(); return tot/n
+net.train(); print("val loss before",evalval(),flush=True)
+t0=time.time(); ptr=0; ema=None
+for step in range(1,STEPS+1):
+    lr=LR*min(1,step/50)*(0.1+0.9*0.5*(1+math.cos(math.pi*step/STEPS)))
+    for g in opt.param_groups: g["lr"]=lr
+    if ptr+BS>len(tr): random.shuffle(tr); ptr=0
+    pre,tgt=batch(tr[ptr:ptr+BS]); ptr+=BS
+    with torch.autocast("cuda",dtype=torch.float16): loss,_,_=N.loss_fn(codec,pre,tgt,0.1)
+    opt.zero_grad(set_to_none=True); scaler.scale(loss).backward(); scaler.unscale_(opt)
+    torch.nn.utils.clip_grad_norm_(net.parameters(),1.0); scaler.step(opt); scaler.update()
+    ema=loss.item() if ema is None else 0.98*ema+0.02*loss.item()
+    if step%50==0: print(f"step {step} loss {ema:.4f} lr {lr:.2e} {time.time()-t0:.0f}s",flush=True)
+    if step%1000==0: print("val loss",evalval(),flush=True)
+print("val loss after",evalval(),flush=True)
+net.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]
+res=[]
+for T,K in [(0.7,20),(0.0,1)]:
+    with torch.autocast("cuda",dtype=torch.float16): outs=N.reply(codec,prompts,temperature=T,top_k=K,top_p=0.9,dev=dev)
+    for t,o in zip(tests,outs): res.append({"temp":T,"user":t,"ncm_reply":o}); print(f"[T={T}] User: {t}\n  NCM: {o}",flush=True)
+# speed: one batch of 1 and 8
+import time as _t
+for B in (1,8):
+    torch.cuda.synchronize(); a=_t.time()
+    with torch.autocast("cuda",dtype=torch.float16): N.reply(codec,prompts[:B]*1,dev=dev)
+    torch.cuda.synchronize(); print(f"gen time batch {B}: {_t.time()-a:.3f}s",flush=True)
+json.dump({"src":src,"n_train":len(tr),"steps":STEPS,"samples":res},open("/kaggle/working/report.json","w"),indent=1)
+sd={"model":{k:v.detach().cpu().half() for k,v in net.state_dict().items()}}
+torch.save(sd,"/kaggle/working/generator_chat.pt")
+print("saved",flush=True)

code/ncm_export_kernel.py

@@ -0,0 +1,138 @@
+import os, subprocess, sys, json, random, time, math
+
+def sh(c): print("$",c,flush=True); subprocess.run(c,shell=True,check=True)
+os.makedirs("/kaggle/working/w",exist_ok=True); os.chdir("/kaggle/working/w")
+if not os.path.exists("ncm"): sh("git clone -q https://github.com/HazyResearch/ncm.git")
+sh("pip install -q omegaconf hydra-core regex onnx onnxruntime onnxscript 2>&1 | tail -1")
+os.makedirs("ck",exist_ok=True)
+for f in ["config.yaml","generator.pt","tokenizer.pt"]:
+    if not os.path.exists("ck/"+f): sh(f"curl -sL -o ck/{f} https://huggingface.co/hazyresearch/ncm-checkpoints/resolve/main/xsum/{f}")
+open('ncmchat.py','w').write('"""NCM chat fine-tune: reuse released XSum NCM (frozen multiscale VQ tokenizer + prefix-conditioned generator),\nfine-tune the generator on (conversation prefix -> short reply). Needs HazyResearch/ncm on NCM_ROOT."""\nimport os, sys, random, torch, torch.nn.functional as F\nfrom omegaconf import OmegaConf\nNCM_ROOT = os.environ.get("NCM_ROOT", "/tmp/ncm"); sys.path.insert(0, NCM_ROOT)\nfrom models.next_concept import NextConceptTransformer\nfrom tokenization.vqvae_tokenizer import VQVAETokenizer\n\ndef load(ck, cfg_path=None, dev="cpu", load_gen=True):\n    CFG = OmegaConf.load(cfg_path or os.path.join(ck, "config.yaml"))\n    CFG.tokenizer.vqvae.checkpoint_path = os.path.join(ck, CFG.checkpoints.tokenizer)\n    tk = VQVAETokenizer.from_config(CFG.tokenizer)\n    tk.vqvae = tk.vqvae.to(dev).eval()\n    for p in tk.vqvae.parameters(): p.requires_grad = False\n    net = NextConceptTransformer(vqvae_local=tk.vqvae, config=CFG.model).to(dev)\n    if load_gen:\n        sd = torch.load(os.path.join(ck, CFG.checkpoints.generator), map_location="cpu", weights_only=False)\n        net.load_state_dict(sd["model"], strict=True)\n    return CFG, tk, net\n\ndef fmt_prompt(turns):\n    """turns: list of (role, text) ending with a user turn -> prompt string"""\n    s = "".join(f"{\'User\' if r==\'user\' else \'Assistant\'}: {t}\\n" for r, t in turns)\n    return s + "Assistant:"\n\nclass Codec:\n    def __init__(self, CFG, tk, net):\n        self.tk, self.net, self.P = tk, net, int(CFG.model.prefix_len)\n        self.pre = tk.pre_tokenizer\n        self.EOS, self.PAD = self.pre.eos_token_id, self.pre.pad_token_id\n        self.ppad = int(CFG.model.prefix_pad_token_id)   # pad id used in the generator prefix\n        self.T = tk.vqvae.context_length                  # target length (48)\n        self.hf = self.pre.tokenizer if hasattr(self.pre, "tokenizer") else None\n    def enc(self, text): return self.pre.encode(text)\n    def prefix_ids(self, prompt):\n        ids = self.enc(prompt)[-self.P:]                   # keep most recent context\n        return ids + [self.ppad] * (self.P - len(ids))\n    def target_ids(self, reply):\n        ids = self.enc(" " + reply.strip())[: self.T - 1] + [self.EOS]\n        return ids + [self.EOS] * (self.T - len(ids))\n\n@torch.no_grad()\ndef target_codes(codec, tgt):                              # tgt: LongTensor (B,T) -> idx (B, sum levels)\n    return codec.tk.vqvae.encode(tgt)\n\ndef build_inputs(codec, prefix, idx, corrupt=0.1):\n    net, vq = codec.net, codec.tk.vqvae.vq\n    with torch.no_grad():\n        idc = idx.clone()\n        if corrupt > 0:                                    # robustness: corrupt conditioning codes (as in paper)\n            m = torch.rand_like(idc, dtype=torch.float) < corrupt\n            idc[m] = torch.randint(0, 4096, (int(m.sum()),), device=idc.device)\n        nc = vq.idx_to_ncm_input(idc)                      # cumulative recon fed to levels 1..K-1\n        zp = codec.tk.vqvae.encoder.wte(prefix)[:, :codec.P, :]\n        q0 = net.build_l0_query_input(prefix.shape[0])\n    return torch.cat([zp, q0, nc], dim=1)\n\ndef loss_fn(codec, prefix, tgt, corrupt=0.1):\n    idx = target_codes(codec, tgt)\n    x = build_inputs(codec, prefix, idx, corrupt)\n    pad_mask = prefix == codec.ppad\n    logits = codec.net(x, prefix_pad_mask=None if not hasattr(codec.net, "_prefix_pad_bias") else pad_mask)\n    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]).float(), idx.reshape(-1)), logits, idx\n\n@torch.no_grad()\ndef reply(codec, prompts, temperature=0.8, top_k=20, top_p=0.9, dev="cpu"):\n    net, tk = codec.net, codec.tk\n    pre = torch.tensor([codec.prefix_ids(p) for p in prompts], device=dev)\n    g = net.generate(None, pre, temperature=temperature, top_k=top_k, top_p=top_p)\n    _, zl = tk.vqvae.vq.idx_to_zhats(g, return_zlevels=True)\n    ids = tk.vqvae.decoder(zl[-1]).argmax(-1).tolist()\n    out = []\n    for r in ids:\n        if codec.EOS in r: r = r[: r.index(codec.EOS)]\n        out.append(codec.pre.decode([t for t in r if t not in (codec.EOS, codec.PAD)], skip_special_tokens=True).strip())\n    return out\n')
+sys.path.insert(0,os.getcwd())
+# 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)

code/ncmchat.py

@@ -0,0 +1,76 @@
+"""NCM chat fine-tune: reuse released XSum NCM (frozen multiscale VQ tokenizer + prefix-conditioned generator),
+fine-tune the generator on (conversation prefix -> short reply). Needs HazyResearch/ncm on NCM_ROOT."""
+import os, sys, random, torch, torch.nn.functional as F
+from omegaconf import OmegaConf
+NCM_ROOT = os.environ.get("NCM_ROOT", "/tmp/ncm"); sys.path.insert(0, NCM_ROOT)
+from models.next_concept import NextConceptTransformer
+from tokenization.vqvae_tokenizer import VQVAETokenizer
+
+def load(ck, cfg_path=None, dev="cpu", load_gen=True):
+    CFG = OmegaConf.load(cfg_path or os.path.join(ck, "config.yaml"))
+    CFG.tokenizer.vqvae.checkpoint_path = os.path.join(ck, CFG.checkpoints.tokenizer)
+    tk = VQVAETokenizer.from_config(CFG.tokenizer)
+    tk.vqvae = tk.vqvae.to(dev).eval()
+    for p in tk.vqvae.parameters(): p.requires_grad = False
+    net = NextConceptTransformer(vqvae_local=tk.vqvae, config=CFG.model).to(dev)
+    if load_gen:
+        sd = torch.load(os.path.join(ck, CFG.checkpoints.generator), map_location="cpu", weights_only=False)
+        net.load_state_dict(sd["model"], strict=True)
+    return CFG, tk, net
+
+def fmt_prompt(turns):
+    """turns: list of (role, text) ending with a user turn -> prompt string"""
+    s = "".join(f"{'User' if r=='user' else 'Assistant'}: {t}\n" for r, t in turns)
+    return s + "Assistant:"
+
+class Codec:
+    def __init__(self, CFG, tk, net):
+        self.tk, self.net, self.P = tk, net, int(CFG.model.prefix_len)
+        self.pre = tk.pre_tokenizer
+        self.EOS, self.PAD = self.pre.eos_token_id, self.pre.pad_token_id
+        self.ppad = int(CFG.model.prefix_pad_token_id)   # pad id used in the generator prefix
+        self.T = tk.vqvae.context_length                  # target length (48)
+        self.hf = self.pre.tokenizer if hasattr(self.pre, "tokenizer") else None
+    def enc(self, text): return self.pre.encode(text)
+    def prefix_ids(self, prompt):
+        ids = self.enc(prompt)[-self.P:]                   # keep most recent context
+        return ids + [self.ppad] * (self.P - len(ids))
+    def target_ids(self, reply):
+        ids = self.enc(" " + reply.strip())[: self.T - 1] + [self.EOS]
+        return ids + [self.EOS] * (self.T - len(ids))
+
+@torch.no_grad()
+def target_codes(codec, tgt):                              # tgt: LongTensor (B,T) -> idx (B, sum levels)
+    return codec.tk.vqvae.encode(tgt)
+
+def build_inputs(codec, prefix, idx, corrupt=0.1):
+    net, vq = codec.net, codec.tk.vqvae.vq
+    with torch.no_grad():
+        idc = idx.clone()
+        if corrupt > 0:                                    # robustness: corrupt conditioning codes (as in paper)
+            m = torch.rand_like(idc, dtype=torch.float) < corrupt
+            idc[m] = torch.randint(0, 4096, (int(m.sum()),), device=idc.device)
+        nc = vq.idx_to_ncm_input(idc)                      # cumulative recon fed to levels 1..K-1
+        zp = codec.tk.vqvae.encoder.wte(prefix)[:, :codec.P, :]
+        q0 = net.build_l0_query_input(prefix.shape[0])
+    return torch.cat([zp, q0, nc], dim=1)
+
+def loss_fn(codec, prefix, tgt, corrupt=0.1):
+    idx = target_codes(codec, tgt)
+    x = build_inputs(codec, prefix, idx, corrupt)
+    pad_mask = prefix == codec.ppad
+    logits = codec.net(x, prefix_pad_mask=None if not hasattr(codec.net, "_prefix_pad_bias") else pad_mask)
+    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]).float(), idx.reshape(-1)), logits, idx
+
+@torch.no_grad()
+def reply(codec, prompts, temperature=0.8, top_k=20, top_p=0.9, dev="cpu"):
+    net, tk = codec.net, codec.tk
+    pre = torch.tensor([codec.prefix_ids(p) for p in prompts], device=dev)
+    g = net.generate(None, pre, temperature=temperature, top_k=top_k, top_p=top_p)
+    _, zl = tk.vqvae.vq.idx_to_zhats(g, return_zlevels=True)
+    ids = tk.vqvae.decoder(zl[-1]).argmax(-1).tolist()
+    out = []
+    for r in ids:
+        if codec.EOS in r: r = r[: r.index(codec.EOS)]
+        out.append(codec.pre.decode([t for t in r if t not in (codec.EOS, codec.PAD)], skip_special_tokens=True).strip())
+    return out

generator_chat.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b21acb483deb0814dbaee271f054fd7ebb43019b18321dd28d2de5dd75a87d6
+size 177107827

training_report.json

@@ -0,0 +1,91 @@
+{
+ "src": {
+  "everyday": 8625,
+  "dolly": 10544,
+  "alpaca": 31323
+ },
+ "n_train": 21084,
+ "steps": 4000,
+ "samples": [
+  {
+   "temp": 0.7,
+   "user": "Hi, how are you?",
+   "ncm_reply": "Hello! How can I help you today?"
+  },
+  {
+   "temp": 0.7,
+   "user": "What is the capital of France?",
+   "ncm_reply": "The capital capital of Paris is Paris."
+  },
+  {
+   "temp": 0.7,
+   "user": "Tell me a joke.",
+   "ncm_reply": "Hello! a weather??"
+  },
+  {
+   "temp": 0.7,
+   "user": "What do you like to do for fun?",
+   "ncm_reply": "Enough to make making fun about wishing your fun, and love books and making you look fun!"
+  },
+  {
+   "temp": 0.7,
+   "user": "Can you recommend a good book?",
+   "ncm_reply": "A book is a good book."
+  },
+  {
+   "temp": 0.7,
+   "user": "How do I boil an egg?",
+   "ncm_reply": "When you boil an egg, you might try."
+  },
+  {
+   "temp": 0.7,
+   "user": "I'm feeling sad today.",
+   "ncm_reply": "I am very sorry."
+  },
+  {
+   "temp": 0.7,
+   "user": "What is the meaning of life?",
+   "ncm_reply": "The meaning meaning of life means that."
+  },
+  {
+   "temp": 0.0,
+   "user": "Hi, how are you?",
+   "ncm_reply": "Hello! How can I help you today?"
+  },
+  {
+   "temp": 0.0,
+   "user": "What is the capital of France?",
+   "ncm_reply": "Paris is a capital capital."
+  },
+  {
+   "temp": 0.0,
+   "user": "Tell me a joke.",
+   "ncm_reply": "you don't help you help you today?"
+  },
+  {
+   "temp": 0.0,
+   "user": "What do you like to do for fun?",
+   "ncm_reply": "What fun items do you like like fun, cake, butter, butter, or.."
+  },
+  {
+   "temp": 0.0,
+   "user": "Can you recommend a good book?",
+   "ncm_reply": "Stars is a good book."
+  },
+  {
+   "temp": 0.0,
+   "user": "How do I boil an egg?",
+   "ncm_reply": "You can tap a fruit cake egg."
+  },
+  {
+   "temp": 0.0,
+   "user": "I'm feeling sad today.",
+   "ncm_reply": "I'm feeling sorry.."
+  },
+  {
+   "temp": 0.0,
+   "user": "What is the meaning of life?",
+   "ncm_reply": "Life is a daily meaning of life."
+  }
+ ]
+}