|
|
""" |
|
|
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") |
