''' A visualizer module for Thinc '''
import seaborn
import matplotlib.pyplot as plt


def visualize_attention(x, y, weights, layer='Encoder', self_attn=True):
    '''
        Visualize self/outer attention
        Args:
            x: sentence
            y: sentence
            weights: (nH, nL, nL)
    '''
    def heatmap(x, y, data, ax):
        seaborn.heatmap(data, square=True, xticklabels=y, yticklabels=x, vmin=0.0, vmax=1.0,
                        cbar_kws = dict(use_gridspec=False, location="top"),
                        ax=ax)
    num = min(weights.shape[0], 4)
    fig, axs = plt.subplots(1, num)
    attn_type = 'self attention' if self_attn else 'outer attention'
    fig.suptitle('{} {} for all the heads'.format(layer, attn_type))
    if len(weights.shape) == 3:
        for i in range(num):
            heatmap(x, y, weights[i], axs[i])
    else:
        raise ValueError("Wrong input weights dimensions")
    plt.show()

Add missing visualizer module