Spaces:
Sleeping
Sleeping
| """ | |
| Visualization generation for BLT patching explorer. | |
| """ | |
| import numpy as np | |
| import librosa | |
| import plotly.graph_objects as go | |
| from plotly.subplots import make_subplots | |
| from typing import List | |
| from .audio import load_audio_file, TOKENS_PER_SECOND | |
| def compute_spectrogram(audio_path: str, sr: int | None = None): | |
| """ | |
| Compute mel-spectrogram from audio file. | |
| Args: | |
| audio_path: Path to audio file | |
| sr: Optional sample rate override. When None, uses the file's native rate. | |
| Returns: | |
| times: Time points (seconds) | |
| frequencies: Frequency bins (Hz) | |
| magnitudes: 2D array (freq x time) in dB | |
| """ | |
| file_sr, audio = load_audio_file(audio_path) | |
| sr = sr or file_sr | |
| n_fft = 1024 | |
| hop_length = 64 | |
| n_mels = 256 | |
| S = librosa.feature.melspectrogram( | |
| y=audio, | |
| sr=sr, | |
| n_fft=n_fft, | |
| hop_length=hop_length, | |
| n_mels=n_mels, | |
| fmax=8000 | |
| ) | |
| S_db = librosa.power_to_db(S, ref=np.max) | |
| times = librosa.frames_to_time( | |
| np.arange(S.shape[1]), sr=sr, hop_length=hop_length | |
| ) | |
| frequencies = librosa.mel_frequencies(n_mels=n_mels, fmax=8000) | |
| return times, frequencies, S_db | |
| def create_visualization( | |
| data: dict, | |
| patches: List[List[int]], | |
| metric_name: str, | |
| threshold: float | |
| ): | |
| """ | |
| Create 4-panel synchronized visualization. | |
| Args: | |
| data: Precomputed data dict | |
| patches: Patch boundaries | |
| metric_name: "entropy" or "perplexity" | |
| threshold: Threshold value | |
| Returns: | |
| Plotly figure | |
| """ | |
| import time | |
| start_time = time.time() | |
| print(f"\n[Visualization] Starting for {data['file_id']}") | |
| print(f" Patches: {len(patches)}, Metric: {metric_name}, Threshold: {threshold}") | |
| fig = make_subplots( | |
| rows=3, cols=1, | |
| subplot_titles=( | |
| 'Spectrogram', | |
| f'{metric_name.capitalize()} + Patch Boundaries', | |
| 'Patch Size Distribution' | |
| ), | |
| specs=[ | |
| [{"type": "heatmap"}], | |
| [{"type": "scatter"}], | |
| [{"type": "bar"}] | |
| ], | |
| vertical_spacing=0.1, | |
| row_heights=[0.4, 0.4, 0.2] | |
| ) | |
| # Get metric data | |
| metric = np.array( | |
| data['entropies'] if metric_name == 'entropy' else data['perplexities'] | |
| ) | |
| # Calculate patch boundary times | |
| boundary_times = [] | |
| for patch in patches: | |
| if patch: | |
| boundary_times.append(patch[0] / TOKENS_PER_SECOND) | |
| duration = data['audio_duration_sec'] | |
| # Limit vlines to avoid slowdown | |
| MAX_VLINES = 200 | |
| draw_boundaries = len(boundary_times) <= MAX_VLINES | |
| if draw_boundaries: | |
| print(f" Will draw {len(boundary_times)} patch boundaries") | |
| else: | |
| print(f" Skipping {len(boundary_times)} boundaries (too many, would be slow)") | |
| # 1. Spectrogram | |
| spec_start = time.time() | |
| times, freqs, mags = compute_spectrogram(data['audio_path']) | |
| print(f" Spectrogram computed in {time.time() - spec_start:.2f}s") | |
| fig.add_trace( | |
| go.Heatmap( | |
| x=times, | |
| y=freqs, | |
| z=mags, | |
| colorscale='Viridis', | |
| showscale=True, | |
| colorbar=dict(title="dB", x=1.05) | |
| ), | |
| row=1, col=1 | |
| ) | |
| # Add patch boundaries to spectrogram | |
| if draw_boundaries: | |
| for t in boundary_times: | |
| fig.add_vline( | |
| x=t, line_width=2, line_dash="solid", | |
| line_color="red", opacity=0.7, | |
| row=1, col=1 | |
| ) | |
| fig.update_xaxes(title_text="Time (s)", row=1, col=1, range=[0, duration]) | |
| fig.update_yaxes(title_text="Frequency (Hz)", row=1, col=1) | |
| # 2. Entropy/Perplexity plot | |
| token_positions = np.arange(len(metric)) | |
| fig.add_trace( | |
| go.Scatter( | |
| x=token_positions, | |
| y=metric, | |
| mode='lines+markers', | |
| name=metric_name.capitalize(), | |
| line=dict(color='blue', width=2), | |
| marker=dict(size=3) | |
| ), | |
| row=2, col=1 | |
| ) | |
| # Add threshold line | |
| fig.add_hline( | |
| y=threshold, line_width=2, line_dash="dash", | |
| line_color="red", opacity=0.7, | |
| annotation_text=f"Threshold: {threshold:.2f}", | |
| annotation_position="right", | |
| row=2, col=1 | |
| ) | |
| # Add patch boundaries | |
| if draw_boundaries: | |
| for patch in patches: | |
| if patch: | |
| fig.add_vline( | |
| x=patch[0], line_width=1, line_dash="solid", | |
| line_color="gray", opacity=0.4, | |
| row=2, col=1 | |
| ) | |
| fig.update_xaxes(title_text="Token Position", row=2, col=1, range=[0, len(metric)]) | |
| fig.update_yaxes(title_text=metric_name.capitalize(), row=2, col=1) | |
| # 3. Patch size distribution | |
| sizes = [len(p) for p in patches] | |
| size_counts = {} | |
| for size in sizes: | |
| size_counts[size] = size_counts.get(size, 0) + 1 | |
| sorted_sizes = sorted(size_counts.keys()) | |
| sorted_counts = [size_counts[s] for s in sorted_sizes] | |
| fig.add_trace( | |
| go.Bar( | |
| x=sorted_sizes, | |
| y=sorted_counts, | |
| name='Patch Sizes', | |
| marker_color='steelblue', | |
| hovertemplate='Size: %{x} tokens<br>Count: %{y}<extra></extra>' | |
| ), | |
| row=3, col=1 | |
| ) | |
| fig.update_xaxes(title_text="Patch Size (tokens)", row=3, col=1) | |
| fig.update_yaxes(title_text="Count", row=3, col=1) | |
| # Update layout | |
| fig.update_layout( | |
| height=1200, | |
| showlegend=False, | |
| hovermode='closest', | |
| title_text=f"BLT Patching Visualization - {data['file_id']}" | |
| ) | |
| total_time = time.time() - start_time | |
| print(f" Total visualization time: {total_time:.2f}s") | |
| if not draw_boundaries: | |
| print(f" ⚠️ Note: Patch boundaries not drawn (>{MAX_VLINES} patches)") | |
| return fig | |