File size: 4,431 Bytes
207e447 | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """Block-diffusion inference for this checkpoint.
Usage:
from blockdiff_inference import load, generate
model, tok, mask_id = load("REPO_ID")
print(generate(model, tok, mask_id, "The capital of France is"))
"""
import sys, torch, torch.nn as nn, torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
BLOCK_LEN = 32
CENTERED_CONV = True
SHIFT = True
MASK_TOKEN = "<|mdm_mask|>"
def _mask(attention_mask, B, T, device, window=None):
i = torch.arange(T, device=device)
m = (i[None, :] // BLOCK_LEN) <= (i[:, None] // BLOCK_LEN)
if window is not None:
m &= (i[:, None] - i[None, :]).abs() < window
m = m[None, None].expand(B, 1, T, T).clone()
if attention_mask is not None and attention_mask.dim() == 2:
m &= attention_mask[:, -T:].bool()[:, None, None, :]
return m | torch.eye(T, dtype=torch.bool, device=device)[None, None]
def _mask_fn(sliding):
def fn(*args, **kw):
emb = kw.get("input_embeds", kw.get("inputs_embeds"))
if emb is None:
emb = next(a for a in args if torch.is_tensor(a) and a.dim() == 3)
win = (getattr(kw.get("config"), "sliding_window", None) or 256) if sliding else None
return _mask(kw.get("attention_mask"), emb.shape[0], emb.shape[1], emb.device, win)
return fn
class CenteredConv1d(nn.Module):
def __init__(self, conv):
super().__init__(); self.conv = conv
L = conv.kernel_size[0]; self.lp, self.rp = (L - 1) // 2, L - 1 - (L - 1) // 2
conv.padding = (0,)
@property
def weight(self): return self.conv.weight
@property
def bias(self): return self.conv.bias
def forward(self, x): return self.conv(F.pad(x, (self.lp, self.rp)))
def load(repo_id, device="cuda", dtype=torch.bfloat16):
tok = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo_id, dtype=dtype, trust_remote_code=True,
attn_implementation="sdpa").to(device).eval()
for t in {sys.modules[type(m).__module__] for m in model.modules()}:
for name in dir(t):
lo = name.lower()
if "mask" in lo and "causal" in lo and callable(getattr(t, name)):
setattr(t, name, _mask_fn("sliding" in lo))
for m in model.modules():
if hasattr(m, "is_causal"): m.is_causal = False
if CENTERED_CONV:
for parent in list(model.modules()):
for cn, ch in list(parent.named_children()):
if isinstance(ch, nn.Conv1d) and ch.groups == ch.in_channels and ch.kernel_size[0] > 1:
setattr(parent, cn, CenteredConv1d(ch))
model.config.use_cache = False
return model, tok, tok.convert_tokens_to_ids(MASK_TOKEN)
def _logits(model, x):
lg = model(input_ids=x).logits
return torch.cat([lg[:, :1], lg[:, :-1]], 1) if SHIFT else lg
@torch.inference_mode()
def generate(model, tok, mask_id, prompt, gen_len=256, tau=0.9, temperature=0.0, repetition_penalty=1.1):
x = tok(prompt, return_tensors="pt").input_ids.to(model.device)
P, L, eos = x.shape[1], BLOCK_LEN, tok.eos_token_id
produced, n = 0, (L - P % L if P % L else L)
while produced < gen_len:
x = torch.cat([x, torch.full((1, n), mask_id, device=x.device)], 1)
s, e = x.shape[1] - n, x.shape[1]
for _ in range(L):
is_mask = x[0, s:e] == mask_id
if not is_mask.any(): break
lg = _logits(model, x)[0, s:e].float(); lg[:, mask_id] = -float("inf")
if repetition_penalty != 1.0 and s > P:
seen = x[0, P:s].unique(); v = lg[:, seen]
lg[:, seen] = torch.where(v > 0, v / repetition_penalty, v * repetition_penalty)
pr = F.softmax(lg / temperature if temperature > 0 else lg, -1)
pred = torch.multinomial(pr, 1)[:, 0] if temperature > 0 else pr.argmax(-1)
conf = pr.gather(-1, pred[:, None])[:, 0].masked_fill(~is_mask, -1.0)
sel = (conf >= tau) & is_mask
if not sel.any(): sel = torch.zeros_like(is_mask); sel[conf.argmax()] = True
x[0, s:e] = torch.where(sel, pred, x[0, s:e])
produced += n; n = L
if (x[0, s:e] == eos).any(): break
out = x[0, P:]
hit = (out == eos).nonzero()
if len(hit): out = out[: hit[0, 0]]
return tok.decode(out, skip_special_tokens=True)
|