"""Visualization utilities for Myanmar Ghost project.""" from pathlib import Path from typing import Any, Dict, List, Optional import matplotlib.pyplot as plt import numpy as np import seaborn as sns def plot_training_curves( history: Dict[str, List[float]], metrics: List[str] = None, title: str = "Training Curves", output_path: Optional[str] = None, figsize: tuple = (12, 8), ) -> plt.Figure: """Plot training curves for multiple metrics. Args: history: Dictionary mapping metric names to lists of values metrics: List of metrics to plot (default: all) title: Plot title output_path: Path to save figure figsize: Figure size Returns: Matplotlib figure """ if metrics is None: metrics = list(history.keys()) n_metrics = len(metrics) n_cols = min(2, n_metrics) n_rows = (n_metrics + n_cols - 1) // n_cols fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize) fig.suptitle(title, fontsize=16) if n_metrics == 1: axes = [axes] else: axes = axes.flatten() if hasattr(axes, 'flatten') else axes for i, metric in enumerate(metrics): ax = axes[i] if i < len(axes) else axes[0] if metric in history: values = history[metric] steps = list(range(len(values))) ax.plot(steps, values, marker='o', markersize=3) ax.set_xlabel('Step/Epoch') ax.set_ylabel(metric.capitalize()) ax.set_title(metric.capitalize()) ax.grid(True, alpha=0.3) # Hide unused subplots for i in range(n_metrics, len(axes)): axes[i].set_visible(False) plt.tight_layout() if output_path: Path(output_path).parent.mkdir(parents=True, exist_ok=True) plt.savefig(output_path, dpi=150, bbox_inches='tight') return fig def plot_confusion_matrix( cm: np.ndarray, class_names: List[str], title: str = "Confusion Matrix", output_path: Optional[str] = None, figsize: tuple = (10, 8), normalize: bool = False, ) -> plt.Figure: """Plot confusion matrix. Args: cm: Confusion matrix class_names: Names of classes title: Plot title output_path: Path to save figure figsize: Figure size normalize: Whether to normalize Returns: Matplotlib figure """ if normalize: cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] fig, ax = plt.subplots(figsize=figsize) sns.heatmap( cm, annot=True, fmt='.2f' if normalize else 'd', cmap='Blues', xticklabels=class_names, yticklabels=class_names, ax=ax, ) ax.set_xlabel('Predicted') ax.set_ylabel('True') ax.set_title(title) plt.tight_layout() if output_path: Path(output_path).parent.mkdir(parents=True, exist_ok=True) plt.savefig(output_path, dpi=150, bbox_inches='tight') return fig def plot_label_distribution( labels: List[Any], class_names: Optional[List[str]] = None, title: str = "Label Distribution", output_path: Optional[str] = None, figsize: tuple = (10, 6), ) -> plt.Figure: """Plot distribution of labels. Args: labels: List of labels class_names: Names of classes title: Plot title output_path: Path to save figure figsize: Figure size Returns: Matplotlib figure """ from collections import Counter counts = Counter(labels) if class_names: labels_order = class_names values = [counts.get(l, 0) for l in labels_order] else: labels_order = list(counts.keys()) values = list(counts.values()) fig, ax = plt.subplots(figsize=figsize) bars = ax.bar(labels_order, values, color='steelblue', alpha=0.7) # Add count labels on bars for bar, count in zip(bars, values): height = bar.get_height() ax.text( bar.get_x() + bar.get_width() / 2., height, f'{int(count)}', ha='center', va='bottom', ) ax.set_xlabel('Class') ax.set_ylabel('Count') ax.set_title(title) ax.grid(True, alpha=0.3, axis='y') plt.tight_layout() if output_path: Path(output_path).parent.mkdir(parents=True, exist_ok=True) plt.savefig(output_path, dpi=150, bbox_inches='tight') return fig def plot_attention_weights( attention_weights: np.ndarray, tokens: List[str], title: str = "Attention Weights", output_path: Optional[str] = None, figsize: tuple = (12, 10), ) -> plt.Figure: """Plot attention weights heatmap. Args: attention_weights: Attention weight matrix tokens: List of tokens title: Plot title output_path: Path to save figure figsize: Figure size Returns: Matplotlib figure """ fig, ax = plt.subplots(figsize=figsize) sns.heatmap( attention_weights, xticklabels=tokens, yticklabels=tokens, cmap='viridis', ax=ax, cbar_kw={'label': 'Attention Weight'}, ) ax.set_xlabel('Key Tokens') ax.set_ylabel('Query Tokens') ax.set_title(title) plt.tight_layout() if output_path: Path(output_path).parent.mkdir(parents=True, exist_ok=True) plt.savefig(output_path, dpi=150, bbox_inches='tight') return fig def plot_loss_landscape( losses: np.ndarray, xlabel: str = "x", ylabel: str = "y", title: str = "Loss Landscape", output_path: Optional[str] = None, figsize: tuple = (10, 6), ) -> plt.Figure: """Plot loss landscape. Args: losses: 2D array of loss values xlabel: Label for x-axis ylabel: Label for y-axis title: Plot title output_path: Path to save figure figsize: Figure size Returns: Matplotlib figure """ fig, ax = plt.subplots(figsize=figsize) if losses.ndim == 1: ax.plot(losses) else: sns.heatmap(losses, ax=ax, cmap='viridis') ax.set_xlabel(xlabel) ax.set_ylabel(ylabel) ax.set_title(title) plt.tight_layout() if output_path: Path(output_path).parent.mkdir(parents=True, exist_ok=True) plt.savefig(output_path, dpi=150, bbox_inches='tight') return fig if __name__ == "__main__": print("Visualization utilities loaded") print("Available functions:") print(" - plot_training_curves") print(" - plot_confusion_matrix") print(" - plot_label_distribution") print(" - plot_attention_weights") print(" - plot_loss_landscape")