mirror of https://github.com/razor-ai/soup.git
86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
"""Rich live training dashboard in the terminal."""
|
|
|
|
from typing import Optional
|
|
|
|
from rich.console import Console
|
|
from rich.live import Live
|
|
from rich.panel import Panel
|
|
|
|
from soup_cli.config.schema import SoupConfig
|
|
|
|
console = Console()
|
|
|
|
|
|
class TrainingDisplay:
|
|
"""Live-updating terminal dashboard for training progress."""
|
|
|
|
def __init__(self, config: SoupConfig, device_name: str = ""):
|
|
self.config = config
|
|
self.device_name = device_name
|
|
self.current_step = 0
|
|
self.total_steps = 0
|
|
self.current_epoch = 0
|
|
self.loss = 0.0
|
|
self.lr = 0.0
|
|
self.grad_norm = 0.0
|
|
self.gpu_mem = ""
|
|
self.speed = 0.0
|
|
self._live: Optional[Live] = None
|
|
|
|
def start(self, total_steps: int):
|
|
"""Start the live display."""
|
|
self.total_steps = total_steps
|
|
self._live = Live(self._render(), console=console, refresh_per_second=2)
|
|
self._live.start()
|
|
|
|
def update(self, step: int, epoch: float, loss: float, lr: float, **kwargs):
|
|
"""Update display with new metrics."""
|
|
self.current_step = step
|
|
self.current_epoch = epoch
|
|
self.loss = loss
|
|
self.lr = lr
|
|
self.grad_norm = kwargs.get("grad_norm", 0.0)
|
|
self.speed = kwargs.get("speed", 0.0)
|
|
self.gpu_mem = kwargs.get("gpu_mem", "")
|
|
|
|
if self._live:
|
|
self._live.update(self._render())
|
|
|
|
def stop(self):
|
|
"""Stop the live display."""
|
|
if self._live:
|
|
self._live.stop()
|
|
|
|
def _render(self) -> Panel:
|
|
"""Render the dashboard panel."""
|
|
if self.total_steps > 0:
|
|
progress_pct = self.current_step / self.total_steps * 100
|
|
else:
|
|
progress_pct = 0
|
|
bar_width = 30
|
|
filled = int(bar_width * progress_pct / 100)
|
|
bar = "#" * filled + "-" * (bar_width - filled)
|
|
|
|
epochs = self.config.training.epochs
|
|
epoch_str = f"Epoch {self.current_epoch:.1f}/{epochs}"
|
|
lines = []
|
|
lines.append(f"{epoch_str} [{bar}] {progress_pct:.0f}%")
|
|
lines.append(f"Step: {self.current_step}/{self.total_steps}")
|
|
lines.append(f"Loss: {self.loss:.4f} LR: {self.lr:.2e}")
|
|
|
|
if self.speed > 0:
|
|
lines.append(f"Speed: {self.speed:.2f} it/s")
|
|
if self.gpu_mem:
|
|
lines.append(f"GPU: {self.gpu_mem}")
|
|
if self.grad_norm > 0:
|
|
lines.append(f"Grad: {self.grad_norm:.4f}")
|
|
|
|
content = "\n".join(lines)
|
|
name = self.config.experiment_name or self.config.base
|
|
return Panel(
|
|
content,
|
|
title=f"[bold green]Soup Training: {name}[/]",
|
|
subtitle=f"[dim]{self.device_name}[/]",
|
|
border_style="green",
|
|
)
|