Plot Attention

plot_attention.py

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

def plot_attention(attention, spa, eng):
    spa = standardize(spa).numpy().decode().split()
    eng = standardize(eng).numpy().decode().split()[1:]
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(1, 1, 1)

    attention = tf.squeeze(attention).numpy()

    ax.matshow(attention, cmap='viridis', vmin=0.0)

    fontdict = {'fontsize': 14}

    ax.set_xticklabels([''] + spa, fontdict=fontdict, rotation=90)
    ax.set_yticklabels([''] + eng, fontdict=fontdict)

    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))

    ax.set_xlabel('Input text')
    ax.set_ylabel('Output text')
    plt.suptitle('Attention weights')