Spaces:
Sleeping
Sleeping
| """Summary plots logged as MLflow artifacts in the transform stage.""" | |
| import json | |
| import logging | |
| import os | |
| import matplotlib.pyplot as plt | |
| import mlflow | |
| import numpy as np | |
| from datasets import load_from_disk | |
| from matplotlib.colors import LogNorm | |
| def _load_refined(refined_path): | |
| df = load_from_disk(refined_path).to_pandas() | |
| with open(os.path.join(refined_path, "metadata.json")) as f: | |
| metadata = json.load(f) | |
| return df, metadata | |
| def _plot_p_w_single(df, w_bins, model_name): | |
| centers = 0.5 * (w_bins[:-1] + w_bins[1:]) | |
| width = w_bins[1] - w_bins[0] | |
| row = df.query("layer == 0 and head == 0").iloc[0] | |
| p_w = np.asarray(row["P_w"], dtype=float) | |
| fig, ax = plt.subplots(figsize=(7, 4)) | |
| ax.bar(centers, p_w, width=width, align="center") | |
| ax.set_xlabel("Weight") | |
| ax.set_ylabel("Probability") | |
| ax.set_title(f"{model_name} L0 H0 — P(W) [W_QK]") | |
| fig.tight_layout() | |
| return fig | |
| def _plot_p_w_stacked(df_sorted, w_bins, n_layers, n_heads, model_name): | |
| Z = np.array([np.asarray(row["P_w"], dtype=float) | |
| for _, row in df_sorted.iterrows()]).T | |
| floor = max(Z[Z > 0].min() * 0.5, 1e-12) if np.any(Z > 0) else 1e-12 | |
| Z_plot = np.maximum(Z, floor) | |
| fig, ax = plt.subplots(figsize=(9, 5)) | |
| im = ax.imshow( | |
| Z_plot, | |
| aspect="auto", | |
| origin="lower", | |
| norm=LogNorm(vmin=floor, vmax=Z_plot.max()), | |
| extent=[0, n_layers * n_heads, w_bins[0], w_bins[-1]], | |
| cmap="viridis", | |
| ) | |
| ax.set_xticks([i * n_heads for i in range(n_layers)]) | |
| ax.set_xticklabels([str(i) for i in range(n_layers)]) | |
| ax.set_xlabel("Layer") | |
| ax.set_ylabel("Weight") | |
| ax.set_title(f"{model_name} — P(W) stacked across heads [W_QK]") | |
| cbar = fig.colorbar(im, ax=ax) | |
| cbar.set_label("P(W) (log)") | |
| fig.tight_layout() | |
| return fig | |
| def _plot_std_heatmap(df_sorted, n_layers, n_heads, model_name): | |
| z = df_sorted.pivot(index="layer", columns="head", values="std").values.astype(float) | |
| fig, ax = plt.subplots(figsize=(7, 5)) | |
| im = ax.imshow(z, aspect="auto", origin="lower", cmap="viridis") | |
| ax.set_xticks(range(n_heads)) | |
| ax.set_yticks(range(n_layers)) | |
| ax.set_xlabel("Head") | |
| ax.set_ylabel("Layer") | |
| ax.set_title(f"{model_name} — std(W_QK) per head") | |
| cbar = fig.colorbar(im, ax=ax) | |
| cbar.set_label("std") | |
| fig.tight_layout() | |
| return fig | |
| def log_transform_plots(refined_path, model_name=None): | |
| """Build and log summary plots for the refined dataset to the active MLflow run.""" | |
| if not os.path.isdir(refined_path): | |
| logging.warning("Refined dataset not found at %s; skipping transform plots.", refined_path) | |
| return | |
| df, metadata = _load_refined(refined_path) | |
| df = df.query("weight_type == 'W_QK'") | |
| if df.empty: | |
| logging.warning("No W_QK rows found in %s; skipping transform plots.", refined_path) | |
| return | |
| if model_name is None: | |
| model_name = df["model"].iloc[0] if "model" in df.columns else "" | |
| if "merged" in metadata and model_name in metadata["merged"]: | |
| meta_for_model = metadata["merged"][model_name] | |
| else: | |
| meta_for_model = metadata | |
| w_bins = np.array(meta_for_model["w_bins"]) | |
| df_sorted = df.sort_values(["layer", "head"]).reset_index(drop=True) | |
| n_layers = int(df_sorted["layer"].max()) + 1 | |
| n_heads = int(df_sorted["head"].max()) + 1 | |
| figs = { | |
| "plots/p_w_layer0_head0.png": _plot_p_w_single(df_sorted, w_bins, model_name), | |
| "plots/p_w_stacked.png": _plot_p_w_stacked(df_sorted, w_bins, n_layers, n_heads, model_name), | |
| "plots/std_heatmap.png": _plot_std_heatmap(df_sorted, n_layers, n_heads, model_name), | |
| } | |
| for path, fig in figs.items(): | |
| mlflow.log_figure(fig, path) | |
| plt.close(fig) | |