File size: 4,065 Bytes
951ac76 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """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
|