| 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) |
| |
| random.seed(0) |
| from datasets import load_dataset |
| ex=[] |
| def add_dialog(utts, maxn=6): |
| |
| 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 |
| |
| 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) |
| |
| 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) |
|
|
