| 148 | plt.savefig(f"visualizations/budget_token_k_norms_{step}.png") |
| 149 | |
| 150 | def plot_heatmap_generic(matrix, step, layer_indices, title, cmap="hot"): |
| 151 | sorted_layers = sorted(layer_indices) |
| 152 | |
| 153 | fig, ax = plt.subplots(figsize=(12, len(sorted_layers))) |
| 154 | |
| 155 | im = ax.imshow( |
| 156 | matrix, cmap=cmap, |
| 157 | aspect="auto", interpolation="nearest" |
| 158 | ) |
| 159 | |
| 160 | ax.set_xlabel('Token position', fontsize=12) |
| 161 | ax.set_ylabel('Layer', fontsize=12) |
| 162 | ax.set_title(f'{title} at step {step}', fontsize=14) |
| 163 | |
| 164 | ax.set_yticks(range(len(sorted_layers))) |
| 165 | ax.set_yticklabels(sorted_layers) |
| 166 | |
| 167 | token_len = matrix.shape[-1] |
| 168 | ax.set_xticks(range(token_len)) |
| 169 | ax.set_xticklabels(range(- token_len + 1, 1)) # including current token |
| 170 | |
| 171 | cbar = plt.colorbar(im, ax=ax) |
| 172 | cbar.set_label('Magnitude', rotation=270, labelpad=20) |
| 173 | |
| 174 | ax.set_xticks(np.arange(token_len) - 0.5, minor=True) |
| 175 | ax.set_yticks(np.arange(len(sorted_layers)) - 0.5, minor=True) |
| 176 | |
| 177 | ax.grid(which="minor", color="gray", linestyle="-", linewidth=0.3, alpha=0.3) |
| 178 | |
| 179 | plt.tight_layout() |
| 180 | plt.savefig(f"visualizations/budget_token_heatmap_{step}_{'_'.join(title.split(' '))}.png") |