| """ |
| Patched diffusion sampler (transformers-compatible). |
| AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True) |
| model.generate(**inputs, streamer=...) |
| Defaults: auto steps (confidence threshold), block_size=1024, |
| min_new_tokens=1, max_new_tokens=max length. |
| """ |
| import os, sys, importlib |
| from collections import Counter |
| import torch |
| import torch.nn as nn |
| from huggingface_hub import snapshot_download |
| from transformers import AutoConfig |
|
|
| try: |
| from transformers import Lfm2MoeConfig as _CFG |
| except Exception: |
| _CFG = None |
|
|
|
|
| class DiffusionLFMForCausalLM(nn.Module): |
| config_class = _CFG |
| _auto_class = None |
|
|
| @classmethod |
| def register_for_auto_class(cls, auto_class="AutoModelForCausalLM"): |
| cls._auto_class = auto_class if isinstance(auto_class, str) else auto_class.__name__ |
|
|
| @classmethod |
| def can_generate(cls): |
| return True |
|
|
| def __init__(self, inner, config): |
| super().__init__() |
| self.inner = inner |
| self.config = config |
| self.generation_config = getattr(inner.model, "generation_config", None) |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): |
| config = kwargs.pop("config", None) |
| if config is None or not hasattr(config, "to_dict"): |
| config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True) |
| dtype = kwargs.pop("torch_dtype", None) or kwargs.pop("dtype", None) or "bfloat16" |
| if isinstance(dtype, torch.dtype): |
| dtype = str(dtype).replace("torch.", "") |
| device_map = kwargs.pop("device_map", "auto") |
| revision = kwargs.pop("revision", None) |
|
|
| path = pretrained_model_name_or_path |
| local = path if os.path.isdir(path) else snapshot_download( |
| path, revision=revision, allow_patterns=["*.py", "*.json", "*.txt", "*.model"]) |
| if local not in sys.path: |
| sys.path.insert(0, local) |
| mod = importlib.import_module("modeling_diffusion_lfm") |
| inner = mod.from_pretrained(path, torch_dtype=dtype, device_map=device_map) |
| inner.eval() |
| return cls(inner, config) |
|
|
| @property |
| def device(self): |
| return self.inner.causal_model.device |
|
|
| @property |
| def dtype(self): |
| return next(self.inner.model.parameters()).dtype |
|
|
| def forward(self, *args, **kwargs): |
| return self.inner.causal_model(*args, **kwargs) |
|
|
| def get_input_embeddings(self): |
| return self.inner.model.get_input_embeddings() |
|
|
| def _max_len(self): |
| c = self.inner.model.config |
| return int(getattr(c, "max_position_embeddings", None) or 131072) |
|
|
| @torch.inference_mode() |
| def generate(self, input_ids=None, attention_mask=None, *, inputs=None, |
| max_new_tokens=None, min_new_tokens=1, do_sample=None, temperature=None, |
| repetition_penalty=None, eos_token_id=None, streamer=None, |
| block_size=1024, threshold=0.9, max_steps=None, rep_window=128, |
| max_same=3, loop_ratio=0.35, ctx=256, **kwargs): |
| if input_ids is None: |
| input_ids = inputs |
| if input_ids is None: |
| raise ValueError("input_ids is required") |
| if input_ids.ndim == 1: |
| input_ids = input_ids[None] |
| if input_ids.shape[0] != 1: |
| raise ValueError("only batch size 1 is supported") |
| if attention_mask is not None: |
| input_ids = input_ids[:, attention_mask[0].bool()] |
|
|
| m = self.inner |
| dev = self.device |
| input_ids = input_ids.to(dev) |
| P0 = input_ids.shape[1] |
| if max_new_tokens is None: |
| max_new_tokens = max(1, self._max_len() - P0) |
| PAD = getattr(m.model.config, "pad_token_id", None) |
| EOS = eos_token_id if eos_token_id is not None else m.model.config.eos_token_id |
| eos_set = set(EOS) if isinstance(EOS, (list, tuple)) else {EOS} |
| temp = float(temperature) if temperature and (do_sample or do_sample is None) else 0.0 |
| rep = float(repetition_penalty) if repetition_penalty else 1.3 |
|
|
| if streamer is not None: |
| streamer.put(input_ids.cpu()) |
|
|
| prompt_ids = input_ids |
| generated = [] |
|
|
| while len(generated) < max_new_tokens: |
| count = min(block_size, max_new_tokens - len(generated)) |
| keep = ctx - P0 |
| hist = prompt_ids.new_tensor([generated[-keep:] if keep > 0 else []]) |
| prefix = torch.cat((prompt_ids, hist), 1)[:, -max(ctx, P0):] |
| p = prefix.shape[1] |
| ids = torch.cat((prefix, prefix.new_full((1, count), PAD)), 1) |
| masked = torch.zeros_like(ids, dtype=torch.bool); masked[:, p:] = True |
| valid = torch.ones_like(masked) |
| ban_eos = len(generated) < min_new_tokens |
| steps_cap = max_steps or count |
| per_step = max(1, -(-count // steps_cap)) |
|
|
| while masked.any(): |
| h = m.hidden(ids, masked, valid, prompt_ids.new_tensor([p])) |
| logits = m.model.lm_head(h[masked]).float() |
| if PAD is not None: |
| logits[:, PAD] = -torch.inf |
|
|
| row = ids[0].tolist(); mrow = masked[0].tolist() |
| pos_all = masked[0].nonzero().squeeze(1).tolist() |
| done = [row[i] for i in range(p, len(row)) if not mrow[i]] |
|
|
| if ban_eos: |
| need = min_new_tokens - len(generated) |
| for j, pos in enumerate(pos_all): |
| if pos - p < need: |
| for e in eos_set: |
| logits[j, e] = -torch.inf |
|
|
| cnt = Counter(generated[-rep_window:] + done) |
| if cnt and rep != 1.0: |
| idx = torch.tensor(list(cnt), device=logits.device) |
| pen = torch.tensor([rep ** min(c, 5) for c in cnt.values()], device=logits.device) |
| sel = logits[:, idx] |
| logits[:, idx] = torch.where(sel > 0, sel / pen, sel * pen) |
|
|
| bc = Counter(done) |
| for t, c in bc.items(): |
| if c >= max_same: |
| logits[:, t] = -torch.inf |
| for j, pos in enumerate(pos_all): |
| for nb in (pos - 1, pos + 1): |
| if (p <= nb < len(row) and not mrow[nb]) or nb == p - 1: |
| logits[j, row[nb]] = -torch.inf |
|
|
| probs = (logits / temp if temp else logits).softmax(-1) |
| tokens = torch.multinomial(probs, 1).squeeze(-1) if temp else probs.argmax(-1) |
| conf = probs.gather(1, tokens[:, None]).squeeze(1) |
|
|
| order = conf.argsort(descending=True).tolist() |
| chosen, cpos, used = [], set(), Counter() |
| for j in order: |
| if conf[j] < threshold and len(chosen) >= per_step: |
| break |
| pos = pos_all[j]; t = int(tokens[j]) |
| if pos - 1 in cpos or pos + 1 in cpos or bc[t] + used[t] >= max_same: |
| continue |
| chosen.append(j); cpos.add(pos); used[t] += 1 |
| if not chosen: |
| chosen = order[:1] |
| ch = torch.tensor(chosen, device=tokens.device) |
| positions = torch.tensor([pos_all[j] for j in chosen], device=ids.device) |
| ids[0, positions] = tokens[ch].to(ids.device) |
| masked[0, positions] = False |
|
|
| canvas = ids[0, p:].tolist() |
| stop_at = next((i for i, t in enumerate(canvas) if t in eos_set), None) |
| block = canvas if stop_at is None else canvas[:stop_at] |
| if len(block) >= 8 and len(set(block)) / len(block) < loop_ratio: |
| break |
| generated.extend(block) |
| if streamer is not None and block: |
| streamer.put(torch.tensor(block)) |
| if stop_at is not None: |
| generated.append(canvas[stop_at]) |
| break |
|
|
| if streamer is not None: |
| streamer.end() |
| new = torch.tensor([generated], dtype=input_ids.dtype, device=dev) |
| return torch.cat([input_ids, new], dim=1) |
|
|
