import os import sys import time import argparse import importlib.util import torch import torch.nn.functional as F from transformers import AutoTokenizer # Tracks how many lines the last visualization printed so we can overwrite it _visualize_last_lines = 0 def try_import_infer_base(base_path: str): """Dynamically import `infer-base.py` as a module and return it, or None on failure.""" if not os.path.exists(base_path): return None try: spec = importlib.util.spec_from_file_location("infer_base", base_path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module except Exception as e: print(f"Warning: failed to import {base_path}: {e}") return None def load_finetuned_model(model_path: str, device: str = 'cuda'): """Load a saved fine-tuned model for inference.""" print(f"Loading model from {model_path}...") checkpoint = torch.load(model_path, map_location=device, weights_only=False) config = checkpoint['config'] # Create model model = DiffusionLLM(config) # Load weights state_dict = checkpoint['model_state'] state_dict = {k: v.float() for k, v in state_dict.items()} model.load_state_dict(state_dict) model = model.to(device) model.eval() num_params = sum(p.numel() for p in model.parameters()) / 1e6 print(f"✓ Loaded model: {num_params:.1f}M parameters") # Print training info if available if 'step' in checkpoint: print(f" Trained for {checkpoint['step']} steps") if 'best_val_loss' in checkpoint: print(f" Best validation loss: {checkpoint['best_val_loss']:.4f}") return model, config @torch.no_grad() def generate_block_diffusion( model, tokenizer, prompt: str, steps: int = 32, block_size: int = 32, max_new_tokens: int = 128, device: str = 'cuda', temperature: float = 0.8, top_k: int = 50, top_p: float = 0.9, repetition_penalty: float = 1.2, no_repeat_ngram_size: int = 3, verbose: bool = True, visualize_fn=None, parallel_blocks: int = 1, ): """ Generate text using block diffusion with sampling controls. If `visualize_fn` is provided it will be called as: visualize_fn(tokenizer, context_ids, mask_block, is_masked, config, clear=True) Returns the decoded generated string (including prompt). """ model.eval() # Encode prompt prompt_ids = tokenizer.encode(prompt, return_tensors="pt").to(device) # Get model config config = model.module.config if hasattr(model, 'module') else getattr(model, 'config', None) if hasattr(model, '_orig_mod'): config = model._orig_mod.config if config is None: raise RuntimeError("Could not determine model config") num_blocks = max_new_tokens // block_size parallel_blocks = min(parallel_blocks, num_blocks) if verbose: print(f"Generating {num_blocks} blocks of {block_size} tokens ({max_new_tokens} max_new_tokens)\n") context_ids = prompt_ids all_generated_tokens = set(prompt_ids[0].tolist()) blocks_generated = 0 while blocks_generated < num_blocks: current_parallel = min(parallel_blocks, num_blocks - blocks_generated) if current_parallel > 1: new_blocks = _generate_parallel_blocks( model, tokenizer, context_ids, config, device, current_parallel, block_size, steps, temperature, top_k, top_p, repetition_penalty, no_repeat_ngram_size, all_generated_tokens, visualize_fn ) for block in new_blocks: context_ids = torch.cat([context_ids, block], dim=1) blocks_generated += 1 else: mask_block, block_token_history = _generate_single_block( model, tokenizer, context_ids, config, device, block_size, steps, temperature, top_k, top_p, repetition_penalty, no_repeat_ngram_size, all_generated_tokens, visualize_fn ) context_ids = torch.cat([context_ids, mask_block], dim=1) blocks_generated += 1 generated_ids = context_ids[0].tolist() return tokenizer.decode(generated_ids, skip_special_tokens=False) def _apply_sampling_controls( block_logits, context_ids, mask_block, is_masked, repetition_penalty, temperature, top_k, top_p, no_repeat_ngram_size, block_token_history ): """Apply repetition penalty, temperature, top-k, top-p, and n-gram blocking.""" if repetition_penalty != 1.0: seen_tokens = set(context_ids[0].tolist()) for i in range(mask_block.shape[1]): if not is_masked[0, i]: seen_tokens.add(mask_block[0, i].item()) for token_id in seen_tokens: if token_id < block_logits.shape[-1]: avg = block_logits[0, :, token_id].mean() if avg > 0: block_logits[:, :, token_id] /= repetition_penalty else: block_logits[:, :, token_id] *= repetition_penalty block_logits = block_logits / temperature if top_k > 0: k = min(top_k, block_logits.size(-1)) top_k_logits, top_k_indices = torch.topk(block_logits, k, dim=-1) filtered = torch.full_like(block_logits, float('-inf')) filtered.scatter_(-1, top_k_indices, top_k_logits) block_logits = filtered if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(block_logits, descending=True, dim=-1) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = 0 indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove) block_logits[indices_to_remove] = float('-inf') if no_repeat_ngram_size > 0 and len(block_token_history) >= no_repeat_ngram_size - 1: recent_ngram = tuple(block_token_history[-(no_repeat_ngram_size - 1):]) full_history = context_ids[0].tolist() + block_token_history for i in range(len(full_history) - no_repeat_ngram_size + 1): if tuple(full_history[i:i + no_repeat_ngram_size - 1]) == recent_ngram: blocked_token = full_history[i + no_repeat_ngram_size - 1] if blocked_token < block_logits.shape[-1]: block_logits[:, :, blocked_token] = float('-inf') # Safety: reset if all logits are -inf all_inf_mask = torch.isinf(block_logits).all(dim=-1) if all_inf_mask.any(): block_logits[all_inf_mask] = 0.0 return block_logits def _generate_single_block( model, tokenizer, context_ids, config, device, block_size, steps, temperature, top_k, top_p, repetition_penalty, no_repeat_ngram_size, all_generated_tokens, visualize_fn=None ): """Generate a single block using diffusion.""" mask_block = torch.full((1, block_size), config.mask_token_id, device=device) is_masked = torch.ones(1, block_size, dtype=torch.bool, device=device) block_token_history = [] for step_idx in range(steps): full_input = torch.cat([context_ids, mask_block], dim=1) attention_mask = torch.ones_like(full_input, dtype=torch.float32) logits, _ = model(full_input, attention_mask=attention_mask) block_logits = logits[:, -block_size:, :] block_logits = _apply_sampling_controls( block_logits, context_ids, mask_block, is_masked, repetition_penalty, temperature, top_k, top_p, no_repeat_ngram_size, block_token_history ) probs = F.softmax(block_logits, dim=-1) probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0) probs = probs.clamp(min=1e-10) probs = probs / probs.sum(dim=-1, keepdim=True) sampled_tokens = torch.multinomial(probs.view(-1, probs.size(-1)), num_samples=1) sampled_tokens = sampled_tokens.view(1, block_size) confidence = probs.gather(-1, sampled_tokens.unsqueeze(-1)).squeeze(-1) tokens_to_unmask = max(1, block_size // steps) if step_idx == steps - 1: tokens_to_unmask = int(is_masked.sum().item()) if tokens_to_unmask > 0 and is_masked.sum() > 0: masked_confidence = confidence.clone() masked_confidence[~is_masked] = -1.0 num_to_unmask = min(int(tokens_to_unmask), int(is_masked.sum().item())) _, top_indices = torch.topk(masked_confidence.view(-1), num_to_unmask) for idx in top_indices: idx = int(idx.item()) mask_block[0, idx] = sampled_tokens[0, idx] is_masked[0, idx] = False block_token_history.append(sampled_tokens[0, idx].item()) all_generated_tokens.add(sampled_tokens[0, idx].item()) if callable(visualize_fn): try: visualize_fn(tokenizer, context_ids, mask_block, is_masked, config, clear=(step_idx > 0)) except Exception: pass elif visualize_fn: visualize_diffusion_state_local(tokenizer, context_ids, mask_block, is_masked, config, clear=(step_idx > 0)) return mask_block, block_token_history def _generate_parallel_blocks( model, tokenizer, context_ids, config, device, num_parallel, block_size, steps, temperature, top_k, top_p, repetition_penalty, no_repeat_ngram_size, all_generated_tokens, visualize_fn=None ): """Generate multiple blocks in parallel using batched computation. Each block sees all previous blocks in the sequence, maintaining proper order: - Block 0: context + [block0] - Block 1: context + [block0] + [block1] - Block 2: context + [block0] + [block1] + [block2] - etc. This ensures sequential coherence while still benefiting from batched computation. """ batch_size = num_parallel context_len = context_ids.shape[1] # Initialize mask blocks for all parallel blocks # Shape: (num_parallel, block_size) mask_blocks = torch.full((batch_size, block_size), config.mask_token_id, device=device) is_masked = torch.ones(batch_size, block_size, dtype=torch.bool, device=device) block_token_histories = [[] for _ in range(batch_size)] for step_idx in range(steps): # Build inputs with proper sequential structure # Each batch item has context + all previous blocks + its own block # Block i sees: context + block_0 + block_1 + ... + block_i # Create padded inputs - each batch item has different length # We'll pad to the longest sequence (which is the last block) max_seq_len = context_len + (num_parallel * block_size) # Build full input for each batch item full_inputs = [] attention_masks = [] for b in range(batch_size): # This block sees: context + all previous blocks + its own block seq_parts = [context_ids[0]] # Start with context # Add all blocks from 0 to b (inclusive) for prev_b in range(b + 1): seq_parts.append(mask_blocks[prev_b]) # Concatenate to form this batch item's input batch_input = torch.cat(seq_parts, dim=0) # (seq_len,) current_len = batch_input.shape[0] # Pad to max_seq_len padding_needed = max_seq_len - current_len if padding_needed > 0: pad_token = config.pad_token_id if config.pad_token_id is not None else 0 padding = torch.full((padding_needed,), pad_token, device=device) batch_input = torch.cat([batch_input, padding], dim=0) full_inputs.append(batch_input) # Create attention mask (1 for real tokens, 0 for padding) attn_mask = torch.zeros(max_seq_len, device=device) attn_mask[:current_len] = 1.0 attention_masks.append(attn_mask) # Stack into batched tensors full_input = torch.stack(full_inputs, dim=0) # (batch, max_seq_len) attention_mask = torch.stack(attention_masks, dim=0) # (batch, max_seq_len) # Single forward pass for all blocks logits, _ = model(full_input, attention_mask=attention_mask) # Extract logits for each block's position # Block b's logits are at positions [context_len + b*block_size : context_len + (b+1)*block_size] block_logits_list = [] for b in range(batch_size): start_pos = context_len + (b * block_size) end_pos = start_pos + block_size block_logits_list.append(logits[b, start_pos:end_pos, :]) block_logits = torch.stack(block_logits_list, dim=0) # (batch, block_size, vocab) # Apply sampling controls per batch item for b in range(batch_size): # Build context that includes previous blocks for repetition penalty extended_context = context_ids if b > 0: prev_blocks = mask_blocks[:b] extended_context = torch.cat([context_ids] + [prev_blocks.view(1, -1)], dim=1) block_logits[b:b+1] = _apply_sampling_controls( block_logits[b:b+1], extended_context, mask_blocks[b:b+1], is_masked[b:b+1], repetition_penalty, temperature, top_k, top_p, no_repeat_ngram_size, block_token_histories[b] ) probs = F.softmax(block_logits, dim=-1) probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0) probs = probs.clamp(min=1e-10) probs = probs / probs.sum(dim=-1, keepdim=True) # Sample for all batches sampled_tokens = torch.multinomial(probs.view(-1, probs.size(-1)), num_samples=1) sampled_tokens = sampled_tokens.view(batch_size, block_size) confidence = probs.gather(-1, sampled_tokens.unsqueeze(-1)).squeeze(-1) tokens_to_unmask = max(1, block_size // steps) if step_idx == steps - 1: tokens_to_unmask = block_size # Unmask all remaining # Unmask for each batch item for b in range(batch_size): if is_masked[b].sum() > 0: masked_confidence = confidence[b] masked_confidence = masked_confidence.clone() masked_confidence[~is_masked[b]] = -1.0 num_to_unmask = min(int(tokens_to_unmask), int(is_masked[b].sum().item())) _, top_indices = torch.topk(masked_confidence.view(-1), num_to_unmask) for idx in top_indices: idx = int(idx.item()) mask_blocks[b, idx] = sampled_tokens[b, idx] is_masked[b, idx] = False block_token_histories[b].append(sampled_tokens[b, idx].item()) all_generated_tokens.add(sampled_tokens[b, idx].item()) if callable(visualize_fn): try: block_list = [mask_blocks[b:b+1] for b in range(batch_size)] is_masked_list = [is_masked[b:b+1] for b in range(batch_size)] visualize_fn(tokenizer, context_ids, block_list, is_masked_list, config, clear=(step_idx > 0)) except Exception: pass elif visualize_fn: block_list = [mask_blocks[b:b+1] for b in range(batch_size)] is_masked_list = [is_masked[b:b+1] for b in range(batch_size)] visualize_diffusion_state_local(tokenizer, context_ids, block_list, is_masked_list, config, clear=(step_idx > 0)) # Return list of generated blocks return [mask_blocks[b:b+1] for b in range(batch_size)] def chat(model, tokenizer, instruction: str, parallel_blocks: int = 1, **kwargs): """Simple chat interface.""" device = next(model.parameters()).device prompt = format_instruct_prompt(instruction) generated = generate_block_diffusion( model, tokenizer, prompt=prompt, device=device, parallel_blocks=parallel_blocks, **kwargs ) # Extract all assistant responses using ChatML tags start_tag = "<|im_start|>assistant" end_tag = "<|im_end|>" resp_parts = [] pos = 0 while True: start_idx = generated.find(start_tag, pos) if start_idx == -1: break start_idx += len(start_tag) end_idx = generated.find(end_tag, start_idx) if end_idx == -1: resp_parts.append(generated[start_idx:].strip()) break resp_parts.append(generated[start_idx:end_idx].strip()) pos = end_idx + len(end_tag) if resp_parts: resp = "\n\n".join(p for p in resp_parts if p) else: # Fallback if no assistant tags found resp = generated.replace("<|im_start|>assistant", "").replace("<|im_end|>", "").strip() return generated, resp def format_instruct_prompt(instruction: str) -> str: """Format instruction using a simple ChatML-like template.""" return ( "<|im_start|>system\n" "Answer this question truthfully<|im_end|>\n" "<|im_start|>user\n" f"{instruction}\n" "<|im_end|>\n" "<|im_start|>assistant\n" ) def visualize_diffusion_state_local(tokenizer, context_ids, mask_blocks, is_masked_list, config, clear=True, block_colors=None): """Local visualization copied from infer-base.py to ensure consistent terminal output.""" import sys import os # Default colors for different blocks (green, cyan, yellow, magenta) DEFAULT_COLORS = ['\033[92m', '\033[96m', '\033[93m', '\033[95m'] MASK_COLOR = '\033[90m' # Gray for masked tokens RESET = '\033[0m' # Normalize inputs to lists if not isinstance(mask_blocks, list): mask_blocks = [mask_blocks] is_masked_list = [is_masked_list] if block_colors is None: block_colors = DEFAULT_COLORS # Decode context (prompt + previously generated blocks) and replace newlines try: context_text = tokenizer.decode(context_ids[0], skip_special_tokens=True).replace('\n', ' ') except Exception: # Fallback to str context_text = str(context_ids[0].tolist()) # Build visualization for all blocks all_blocks_text = [] for block_idx, (mask_block, is_masked) in enumerate(zip(mask_blocks, is_masked_list)): color = block_colors[block_idx % len(block_colors)] block_tokens = mask_block[0].tolist() block_color_tokens = [] for i, token_id in enumerate(block_tokens): if is_masked[0, i]: # Use block-specific color for masked tokens to distinguish blocks block_color_tokens.append(f'{MASK_COLOR}██{RESET}') else: # Decode individual token; use block color for revealed tokens try: token_text = tokenizer.decode([token_id], skip_special_tokens=False) except Exception: token_text = str(int(token_id)) block_color_tokens.append(f'{color}{token_text}{RESET}') all_blocks_text.append(''.join(block_color_tokens)) # Join all blocks with a subtle separator blocks_combined = ''.join(all_blocks_text) # Overwrite previous visualization area (if any) by moving cursor up and clearing lines. # This prevents accumulation of repeated frames in terminals like VSCode integrated terminal. global _visualize_last_lines if clear and _visualize_last_lines > 0: try: # Move cursor up to the start of the previous block sys.stdout.write(f'\x1b[{_visualize_last_lines}A') # Clear each line that was previously printed for _ in range(_visualize_last_lines): sys.stdout.write('\x1b[2K') # Erase entire line sys.stdout.write('\x1b[1B') # Move cursor down one line # Move cursor back to the top of cleared region sys.stdout.write(f'\x1b[{_visualize_last_lines}A') sys.stdout.flush() except Exception: # Fallback to whole-screen clear try: sys.stdout.write('\x1b[2J\x1b[H') sys.stdout.flush() except Exception: try: clear_cmd = 'cls' if os.name == 'nt' else 'clear' os.system(clear_cmd) except Exception: sys.stdout.write('\r\033[K') sys.stdout.flush() elif clear: # No previous region to overwrite; do a simple ANSI clear to start fresh try: sys.stdout.write('\x1b[2J\x1b[H') sys.stdout.flush() except Exception: try: clear_cmd = 'cls' if os.name == 'nt' else 'clear' os.system(clear_cmd) except Exception: sys.stdout.write('\r\033[K') sys.stdout.flush() # Print legend for parallel blocks if len(mask_blocks) > 1: legend_parts = [] for i in range(len(mask_blocks)): color = block_colors[i % len(block_colors)] legend_parts.append(f'{color}Block {i+1}{RESET}') print(f"Generating: {' | '.join(legend_parts)}\n") # Print the full context with colored blocks # Ensure trailing newline so subsequent clears have predictable behavior out_text = f"{context_text}{blocks_combined}\n" try: sys.stdout.write(out_text) sys.stdout.flush() except Exception: print(out_text, flush=True) # Update last-lines counter so next frame can overwrite this one try: _visualize_last_lines = out_text.count('\n') + (1 if len(mask_blocks) > 1 else 0) + 1 except Exception: _visualize_last_lines = out_text.count('\n') def main(): base_path = os.path.join(os.path.dirname(__file__), "infer-base.py") base_mod = try_import_infer_base(base_path) if base_mod is None or not hasattr(base_mod, 'DiffusionLLM'): raise RuntimeError("DiffusionLLM not found in infer-base.py") DiffusionLLM = base_mod.DiffusionLLM # Workaround for torch.load pickling try: main_mod = sys.modules.get('__main__') if main_mod is not None: if hasattr(base_mod, 'ModelConfig'): setattr(main_mod, 'ModelConfig', base_mod.ModelConfig) setattr(main_mod, 'DiffusionLLM', DiffusionLLM) except Exception: pass parser = argparse.ArgumentParser() parser.add_argument("--model", type=str, default="./checkpoints/model_fp32.pt", help="Path to model checkpoint") parser.add_argument("--tokenizer", type=str, default="Qwen/Qwen2.5-0.5B", help="Tokenizer model id or path") parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu") parser.add_argument("--visualize", action="store_true", default=False, help="Enable visualization during generation") parser.add_argument("--steps", type=int, default=64) parser.add_argument("--block_size", type=int, default=128) parser.add_argument("--max_new_tokens", type=int, default=128) parser.add_argument("--parallel_blocks", type=int, default=1, help="Number of blocks to generate in parallel") args = parser.parse_args() device = torch.device(args.device) print(f"Using device: {device}") # Load tokenizer print("Loading tokenizer...") tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) if tokenizer.pad_token is None: # set pad token if not present tokenizer.pad_token = tokenizer.eos_token # Load model best_model_path = "checkpoints/best_model.pt" if os.path.exists(best_model_path): print("Loading best model...") model, config = load_finetuned_model(best_model_path, device) else: model, config = load_finetuned_model(args.model, device) # Use the local visualization implementation for consistency visualize_fn = None if args.visualize: visualize_fn = visualize_diffusion_state_local print("Ready. Type a message and press Enter (empty line to quit).\n") while True: try: user_input = input("User: ").strip() except (EOFError, KeyboardInterrupt): print("\nExiting.") break if user_input == "": print("Goodbye.") break raw_output, response = chat( model, tokenizer, user_input, steps=args.steps, block_size=args.block_size, max_new_tokens=args.max_new_tokens, temperature=0.8, top_k=50, top_p=0.9, repetition_penalty=1.2, no_repeat_ngram_size=3, verbose=False, visualize_fn=visualize_fn, parallel_blocks=args.parallel_blocks, ) print("\nRaw Output:\n") print(raw_output) print("\nAssistant:\n") print(response) print("\n" + ("=" * 60) + "\n") if __name__ == "__main__": main()