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)