| """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) |
| self.T = tk.vqvae.context_length |
| 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:] |
| 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): |
| 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: |
| 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) |
| 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 |
|
|
