qoherent/ria-toolkit-oss
19
6
Fork 0
Code Issues Pull Requests 1 Workflows Packages Projects Releases 6 Wiki Activity

Widget Support Panels #5

Merged
benchinnery merged 9 commits from new_widgets_support into main 2025-10-22 12:04:53 -04:00
Conversation 0 Commits 9 Files Changed 3 +127 -132
1 changed files with 127 additions and 132 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit c06e58f5d6 - Show all commits

253
src/ria_toolkit_oss/viz/pytorch_state_dict.py
View File

@ -1,188 +1,190 @@
import numpy as np
import plotly.graph_objects as go
from plotly.graph_objects import Figure
import numpy as np
def create_styled_error_figure(title: str, message: str, suggestion: str = None) -> go.Figure:
"""Create a professional error figure with Qoherent dark theme styling."""
fig = go.Figure()
# Create a clean, centered text display using Plotly's text formatting
main_text = f"<b style='color:#f56565;font-size:18px'>⚠️ {title}</b><br><br>"
main_text += f"<span style='color:#e2e8f0;font-size:14px'>{message}</span>"
if suggestion:
main_text += "<br><br><span style='color:#63b3ed;font-size:13px'>💡 <b>Suggestion:</b></span><br>"
main_text += f"<span style='color:#cbd5e0;font-size:12px'>{suggestion}</span>"
# Add the main text annotation
fig.add_annotation(
text=main_text,
xref="paper",
yref="paper",
x=0.5,
y=0.5,
xanchor="center",
yanchor="middle",
showarrow=False,
align="center",
borderwidth=2,
bordercolor="#4a5568",
bgcolor="#2d3748",
font=dict(family="Arial, sans-serif", size=14, color="#e2e8f0"),
)
# Update layout with dark theme
fig.update_layout(
title="",
height=400,
template="plotly_dark",
margin=dict(l=40, r=40, t=40, b=40),
plot_bgcolor="#1a202c",
paper_bgcolor="#1a202c",
font=dict(color="#e2e8f0"),
)
# Remove axes and grid
fig.update_xaxes(visible=False)
fig.update_yaxes(visible=False)
return fig
def model_summary_plot(state_dict: dict) -> Figure:
"""Generate a summary plot of the PyTorch model state dict."""
if not state_dict:
# Handle empty state dict
fig = go.Figure()
fig.add_annotation(
text="No parameters found in state dict",
xref="paper",
yref="paper",
x=0.5,
y=0.5,
showarrow=False,
font=dict(size=16),
return create_styled_error_figure(
"Empty State Dict",
"No parameters found in state dict",
"Ensure the model state dictionary contains weight parameters"
)
fig.update_layout(
title="Model Layer Parameter Counts",
xaxis_title="Layer",
yaxis_title="Number of Parameters",
template="plotly_dark",
)
return fig
# Count parameters by layer type
layer_info = []
for key, tensor in state_dict.items():
if "weight" in key:
if 'weight' in key:
try:
layer_name = key.replace(".weight", "")
layer_name = key.replace('.weight', '')
param_count = (
tensor.numel()
if hasattr(tensor, "numel")
else len(tensor.flatten()) if hasattr(tensor, "flatten") else 0
tensor.numel() if hasattr(tensor, 'numel')
else len(tensor.flatten()) if hasattr(tensor, 'flatten')
else 0
)
shape = (
list(tensor.shape)
if hasattr(tensor, "shape")
else [len(tensor)] if hasattr(tensor, "__len__") else []
list(tensor.shape) if hasattr(tensor, 'shape')
else [len(tensor)] if hasattr(tensor, '__len__')
else []
)
layer_info.append({"layer": layer_name, "parameters": param_count, "shape": shape})
layer_info.append({
'layer': layer_name,
'parameters': param_count,
'shape': shape
})
except Exception as e:
print(f"Warning: Could not process layer {key}: {e}")
continue
if not layer_info:
# Handle case where no weight layers found
fig = go.Figure()
fig.add_annotation(
text="No weight layers found in state dict",
xref="paper",
yref="paper",
x=0.5,
y=0.5,
showarrow=False,
font=dict(size=16),
return create_styled_error_figure(
"No Weight Layers Found",
"No weight layers found in state dict",
"Ensure the state dictionary contains layers with '.weight' parameters"
)
fig.update_layout(
title="Model Layer Parameter Counts",
xaxis_title="Layer",
yaxis_title="Number of Parameters",
template="plotly_dark",
)
return fig
# Create bar chart of parameter counts
fig = go.Figure(
data=[
fig = go.Figure(data=[
go.Bar(
x=[info["layer"] for info in layer_info],
y=[info["parameters"] for info in layer_info],
x=[info['layer'] for info in layer_info],
y=[info['parameters'] for info in layer_info],
text=[f"Shape: {info['shape']}" for info in layer_info],
textposition="auto",
textposition='auto',
)
]
)
])
fig.update_layout(
title="Model Layer Parameter Counts",
xaxis_title="Layer",
yaxis_title="Number of Parameters",
template="plotly_dark",
template="plotly_dark"
)
return fig
def layer_weights_plot(state_dict: dict, layer_name: str = None) -> Figure:
"""Visualize weights for a specific layer."""
if not state_dict:
fig = go.Figure()
fig.add_annotation(
text="No data in state dict", xref="paper", yref="paper", x=0.5, y=0.5, showarrow=False, font=dict(size=16)
return create_styled_error_figure(
"Empty State Dict",
"No data in state dict",
"Ensure the model state dictionary contains data"
)
fig.update_layout(title="Layer Weights", template="plotly_dark")
return fig
if layer_name is None:
# Get first weight tensor
weight_keys = [k for k in state_dict.keys() if "weight" in k]
weight_keys = [k for k in state_dict.keys() if 'weight' in k]
if not weight_keys:
fig = go.Figure()
fig.add_annotation(
text="No weight tensors found in state dict",
xref="paper",
yref="paper",
x=0.5,
y=0.5,
showarrow=False,
font=dict(size=16),
return create_styled_error_figure(
"No Weight Tensors Found",
"No weight tensors found in state dict",
"Ensure the state dictionary contains layers with '.weight' parameters"
)
fig.update_layout(title="Layer Weights", template="plotly_dark")
return fig
layer_name = weight_keys[0]
try:
weights = state_dict[layer_name]
# Convert to numpy if it's a torch tensor
if hasattr(weights, "numpy"):
weights_np = weights.detach().numpy() if hasattr(weights, "detach") else weights.numpy()
elif hasattr(weights, "cpu"):
if hasattr(weights, 'numpy'):
weights_np = weights.detach().numpy() if hasattr(weights, 'detach') else weights.numpy()
elif hasattr(weights, 'cpu'):
weights_np = weights.cpu().detach().numpy()
else:
weights_np = np.array(weights)
# For 2D weights, create heatmap
if len(weights_np.shape) == 2:
fig = go.Figure(data=go.Heatmap(z=weights_np, colorscale="RdBu", zmid=0))
fig.update_layout(title=f"Weights Heatmap: {layer_name}", template="plotly_dark")
fig = go.Figure(data=go.Heatmap(
z=weights_np,
colorscale='RdBu',
zmid=0
))
fig.update_layout(
title=f"Weights Heatmap: {layer_name}",
template="plotly_dark"
)
else:
# For other shapes, flatten and show histogram
flat_weights = weights_np.flatten()
fig = go.Figure(data=[go.Histogram(x=flat_weights, nbinsx=50)])
fig.update_layout(title=f"Weight Distribution: {layer_name}", template="plotly_dark")
fig.update_layout(
title=f"Weight Distribution: {layer_name}",
template="plotly_dark"
)
return fig
except Exception as e:
fig = go.Figure()
fig.add_annotation(
text=f"Error processing layer {layer_name}: {str(e)}",
xref="paper",
yref="paper",
x=0.5,
y=0.5,
showarrow=False,
font=dict(size=14),
return create_styled_error_figure(
"Layer Processing Error",
f"Error processing layer {layer_name}: {str(e)}",
"Check that the layer name exists and contains valid tensor data"
)
fig.update_layout(title="Layer Weights - Error", template="plotly_dark")
return fig
def weight_distribution_plot(state_dict: dict) -> Figure:
"""Show distribution of weights across all layers."""
if not state_dict:
fig = go.Figure()
fig.add_annotation(
text="No data in state dict", xref="paper", yref="paper", x=0.5, y=0.5, showarrow=False, font=dict(size=16)
return create_styled_error_figure(
"Empty State Dict",
"No data in state dict",
"Ensure the model state dictionary contains data"
)
fig.update_layout(
title="Overall Weight Distribution",
xaxis_title="Weight Value",
yaxis_title="Frequency",
template="plotly_dark",
)
return fig
all_weights = []
layer_names = []
for key, tensor in state_dict.items():
if "weight" in key:
if 'weight' in key:
try:
# Convert to numpy if it's a torch tensor
if hasattr(tensor, "numpy"):
weights_np = tensor.detach().numpy() if hasattr(tensor, "detach") else tensor.numpy()
elif hasattr(tensor, "cpu"):
if hasattr(tensor, 'numpy'):
weights_np = tensor.detach().numpy() if hasattr(tensor, 'detach') else tensor.numpy()
elif hasattr(tensor, 'cpu'):
weights_np = tensor.cpu().detach().numpy()
else:
weights_np = np.array(tensor)
flat_weights = weights_np.flatten()
all_weights.extend(flat_weights)
layer_names.extend([key] * len(flat_weights))
@ -191,31 +193,24 @@ def weight_distribution_plot(state_dict: dict) -> Figure:
continue
if not all_weights:
fig = go.Figure()
fig.add_annotation(
text="No weight data found in state dict",
xref="paper",
yref="paper",
x=0.5,
y=0.5,
showarrow=False,
font=dict(size=16),
return create_styled_error_figure(
"No Weight Data Found",
"No weight data found in state dict",
"Ensure the state dictionary contains layers with '.weight' parameters"
)
fig.update_layout(
title="Overall Weight Distribution",
xaxis_title="Weight Value",
yaxis_title="Frequency",
template="plotly_dark",
)
return fig
fig = go.Figure(data=[go.Histogram(x=all_weights, nbinsx=100, name="All Weights")])
fig = go.Figure(data=[
go.Histogram(
x=all_weights,
nbinsx=100,
name="All Weights"
)
])
fig.update_layout(
title="Overall Weight Distribution",
xaxis_title="Weight Value",
yaxis_title="Frequency",
template="plotly_dark",
template="plotly_dark"
)
return fig