| | """ |
| | Terminal visualization for RND1 generation. |
| | |
| | This module provides real-time visualization of the diffusion denoising process, |
| | showing token evolution and generation progress in the terminal using rich |
| | formatting when available. |
| | """ |
| |
|
| | import torch |
| | from typing import Optional |
| | from tqdm import tqdm |
| |
|
| | try: |
| | from rich.console import Console |
| | from rich.live import Live |
| | from rich.text import Text |
| | from rich.panel import Panel |
| | from rich.progress import Progress, BarColumn, TextColumn, TimeRemainingColumn, MofNCompleteColumn |
| | from rich.layout import Layout |
| | RICH_AVAILABLE = True |
| | except ImportError: |
| | RICH_AVAILABLE = False |
| |
|
| |
|
| | class TerminalVisualizer: |
| | """ |
| | Rich-based visualization for diffusion process with live updates. |
| | |
| | Provides real-time visualization of the token denoising process during |
| | diffusion-based language generation, with colored highlighting of masked |
| | positions and progress tracking. |
| | """ |
| |
|
| | def __init__(self, tokenizer, show_visualization: bool = True): |
| | """ |
| | Initialize the terminal visualizer. |
| | |
| | Args: |
| | tokenizer: The tokenizer for decoding tokens to text |
| | show_visualization: Whether to show visualization (requires rich) |
| | """ |
| | self.tokenizer = tokenizer |
| | self.show_visualization = show_visualization and RICH_AVAILABLE |
| | if not RICH_AVAILABLE and show_visualization: |
| | print("Warning: Install 'rich' for better visualization. Falling back to simple progress bar.") |
| | self.show_visualization = False |
| |
|
| | if self.show_visualization: |
| | self.console = Console() |
| | self.live = None |
| | self.progress = None |
| | self.layout = None |
| | else: |
| | self.pbar = None |
| |
|
| | self.current_tokens = None |
| | self.mask_positions = None |
| | self.total_steps = 0 |
| | self.current_step = 0 |
| |
|
| | def start_visualization(self, initial_tokens: torch.LongTensor, mask_positions: torch.BoolTensor, total_steps: int): |
| | """ |
| | Start the visualization. |
| | |
| | Args: |
| | initial_tokens: Initial token IDs (possibly masked) |
| | mask_positions: Boolean mask indicating which positions are masked |
| | total_steps: Total number of diffusion steps |
| | """ |
| | if not self.show_visualization: |
| | self.pbar = tqdm(total=total_steps, desc="Diffusion") |
| | return |
| |
|
| | self.current_tokens = initial_tokens.clone() |
| | self.mask_positions = mask_positions |
| | self.total_steps = total_steps |
| | self.current_step = 0 |
| |
|
| | self.layout = Layout() |
| | self.layout.split_column( |
| | Layout(name="header", size=3), |
| | Layout(name="text", ratio=1), |
| | Layout(name="progress", size=3) |
| | ) |
| |
|
| | self.progress = Progress( |
| | TextColumn("[bold blue]Diffusion"), |
| | BarColumn(), |
| | MofNCompleteColumn(), |
| | TextColumn("•"), |
| | TextColumn("[cyan]Masks: {task.fields[masks]}"), |
| | TimeRemainingColumn(), |
| | ) |
| | self.progress_task = self.progress.add_task( |
| | "Generating", |
| | total=total_steps, |
| | masks=mask_positions.sum().item() |
| | ) |
| |
|
| | self.live = Live(self.layout, console=self.console, refresh_per_second=4) |
| | self.live.start() |
| | self._update_display() |
| |
|
| | def update_step(self, tokens: torch.LongTensor, maskable: Optional[torch.BoolTensor], step: int, |
| | entropy: Optional[torch.FloatTensor] = None, confidence: Optional[torch.FloatTensor] = None): |
| | """ |
| | Update visualization for current step. |
| | |
| | Args: |
| | tokens: Current token IDs |
| | maskable: Boolean mask of remaining masked positions |
| | step: Current step number |
| | entropy: Optional entropy scores for each position |
| | confidence: Optional confidence scores for each position |
| | """ |
| | if not self.show_visualization: |
| | if self.pbar: |
| | self.pbar.update(1) |
| | masks = maskable.sum().item() if maskable is not None else 0 |
| | self.pbar.set_postfix({'masks': masks}) |
| | return |
| |
|
| | self.current_tokens = tokens.clone() |
| | self.mask_positions = maskable |
| | self.current_step = step |
| |
|
| | masks_remaining = maskable.sum().item() if maskable is not None else 0 |
| | self.progress.update( |
| | self.progress_task, |
| | advance=1, |
| | masks=masks_remaining |
| | ) |
| |
|
| | self._update_display() |
| |
|
| | def _update_display(self): |
| | """Update the live display.""" |
| | if not self.live: |
| | return |
| |
|
| | header = Text("🎭 RND1-Base Generation", style="bold magenta", justify="center") |
| | self.layout["header"].update(Panel(header, border_style="bright_blue")) |
| |
|
| | text_display = self._format_text_with_masks() |
| | self.layout["text"].update( |
| | Panel( |
| | text_display, |
| | title="[bold]Generated Text", |
| | subtitle=f"[dim]Step {self.current_step}/{self.total_steps}[/dim]", |
| | border_style="cyan" |
| | ) |
| | ) |
| |
|
| | self.layout["progress"].update(Panel(self.progress)) |
| |
|
| | def _format_text_with_masks(self) -> Text: |
| | """ |
| | Format text with colored masks. |
| | |
| | Returns: |
| | Rich Text object with formatted tokens |
| | """ |
| | text = Text() |
| |
|
| | if self.current_tokens is None: |
| | return text |
| |
|
| | token_ids = self.current_tokens[0] if self.current_tokens.dim() > 1 else self.current_tokens |
| | mask_flags = self.mask_positions[0] if self.mask_positions is not None and self.mask_positions.dim() > 1 else self.mask_positions |
| |
|
| | for i, token_id in enumerate(token_ids): |
| | if mask_flags is not None and i < len(mask_flags) and mask_flags[i]: |
| | |
| | text.append("[MASK]", style="bold red on yellow" if self.current_step % 2 == 0 else "bold yellow on red") |
| | else: |
| | try: |
| | token_str = self.tokenizer.decode([token_id.item()], skip_special_tokens=False) |
| | |
| | if token_str not in ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<s>", "</s>"]: |
| | |
| | text.append(token_str, style="green" if i < len(token_ids) // 2 else "cyan") |
| | except: |
| | continue |
| |
|
| | return text |
| |
|
| | def stop_visualization(self): |
| | """Stop the visualization and display final result.""" |
| | if not self.show_visualization: |
| | if self.pbar: |
| | self.pbar.close() |
| | print("\n✨ Generation complete!\n") |
| | return |
| |
|
| | if self.live: |
| | self.live.stop() |
| |
|
| | self.console.print("\n[bold green]✨ Generation complete![/bold green]\n") |
| |
|
| | |
| | if self.current_tokens is not None: |
| | try: |
| | token_ids = self.current_tokens[0] if self.current_tokens.dim() > 1 else self.current_tokens |
| | final_text = self.tokenizer.decode(token_ids, skip_special_tokens=True) |
| |
|
| | self.console.print(Panel( |
| | final_text, |
| | title="[bold]Final Generated Text", |
| | border_style="green", |
| | padding=(1, 2) |
| | )) |
| | except: |
| | pass |
| |
|
| |
|
| | class SimpleProgressBar: |
| | """ |
| | Simple progress bar fallback when rich is not available. |
| | |
| | Provides basic progress tracking using tqdm when the rich library |
| | is not installed. |
| | """ |
| |
|
| | def __init__(self, total_steps: int): |
| | """ |
| | Initialize simple progress bar. |
| | |
| | Args: |
| | total_steps: Total number of steps |
| | """ |
| | self.pbar = tqdm(total=total_steps, desc="Diffusion") |
| |
|
| | def update(self, masks_remaining: int = 0): |
| | """ |
| | Update progress bar. |
| | |
| | Args: |
| | masks_remaining: Number of masks still remaining |
| | """ |
| | self.pbar.update(1) |
| | self.pbar.set_postfix({'masks': masks_remaining}) |
| |
|
| | def close(self): |
| | """Close the progress bar.""" |
| | self.pbar.close() |
| | print("\n✨ Generation complete!\n") |
