""" 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: # まだ1トークンも出ていない位置より前で 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)