Move visualization code into AttentionPlotter namespace.