textformer.utils.visualization

Attention-based visualization.

textformer.utils.visualization.plot_attention(input_text, translated_text, attentions, color_map='hot')

Plots an attention graph between input text and translated text.

Parameters

• input_text (list) – List of input tokens.

• translated_text (list) – List of translated tokens.

• attentions (torch.Tensor) – Tensor holding attention values for each (input, translated) token pair.

• color_map (str) – A matplotlib identifier for the color mapping to be used.