Visualizing self attention about easynlp HOT 1 CLOSED

Use the following self-attention shown code from prompt-to-prompt and adapt it to your existing code.
def show_self_attention_comp(attention_store: AttentionStore, res: int, from_where: List[str],
max_com=10, select: int = 0):
attention_maps = aggregate_attention(attention_store, res, from_where, False, select).numpy().reshape((res ** 2, res ** 2))
u, s, vh = np.linalg.svd(attention_maps - np.mean(attention_maps, axis=1, keepdims=True))
images = []
for i in range(max_com):
image = vh[i].reshape(res, res)
image = image - image.min()
image = 255 * image / image.max()
image = np.repeat(np.expand_dims(image, axis=2), 3, axis=2).astype(np.uint8)
image = Image.fromarray(image).resize((256, 256))
image = np.array(image)
images.append(image)
ptp_utils.view_images(np.concatenate(images, axis=1))
from: https://github.com/google/prompt-to-prompt/blob/main/prompt-to-prompt_stable.ipynb

from easynlp.