Transformerモデルのattention可視化

transformerモデルの推論の際のAttentionがどのように作用しているかどうかを可視化して頂きたいです。 HTML出力、画像出力等出力方法は問いません。 イメージとしては以下のようになります。 https://qiita.com/m__k/items/e312ddcf9a3d0ea64d72 モデルは以下になりますのでコード等参照してください。なるべく大きく書き換えることなく実装していただけますと幸いです。 https://github.com/ischlag/TP-Transformer このような依頼を出すのは初めてですので不足事項などありましたらお申し付けください。 以上、ご提案お待ちしております。