Sequence Labelling at paragraph/sentence embedding level using Bi-LSTM + CRF with Keras

I am working on a Sequence tagging task where the element to be tagged is sentences (or paragraph). Most of implementations I found present solution at the token level (NER, POS-Tagging, etc.), whereas here, I first need to build a paragraph embedding before doing sequence labelling.

There is actually two steps: the first model is building a dense representation (with attention) from the word embeddings First model takes the input (#batch size, #words in a paragraph)

the second model takes has input the following shape: - (#batch size, #paragraphs, #words) is taking the output of the first model and apply the Bi-LSTM + CRF architecture on top of it to tag every paragraph of a given document.

import keras 
import keras.backend as K
from keras.layers import *
from keras.activations import *
from keras.regularizers import *
from keras.initializers import *

from keras.models import Model
from keras.layers import Embedding, Input, Bidirectional, GRU, LSTM, TimeDistributed, Layer
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.text import Tokenizer, text_to_word_sequence
from keras_contrib.layers import CRF
from keras_contrib.metrics import crf_accuracy
from keras_contrib.losses import crf_loss
from keras_contrib.metrics import crf_viterbi_accuracy

MAX_PARAGRAPHS = 200 # number of max paragraphs (ie. sequences)
MAX_PARAGRAPH_LENGTH = 60 # number of max token in a given sequences 
EMBEDDING_DIM = 100 # word embedding dimension 
PARAGRAPH_DIM = 100 # paragraph/sentence embedding dimension
MAX_NB_WORDS = 5000 # max nb words in the vocabulary 

n_tags = 5 #number of classes

class AttLayer(Layer):
    def __init__(self, attention_dim):
        self.init = initializers.get('normal')
        self.supports_masking = True
        self.attention_dim = attention_dim
        super(AttLayer, self).__init__()

    def build(self, input_shape):
        assert len(input_shape) == 3
        self.W = K.variable(self.init((input_shape[-1], self.attention_dim)))
        self.b = K.variable(self.init((self.attention_dim, )))
        self.u = K.variable(self.init((self.attention_dim, 1)))
        self.trainable_weights = [self.W, self.b, self.u]
        super(AttLayer, self).build(input_shape)

    def compute_mask(self, inputs, mask=None):
        return mask

    def call(self, x, mask=None):
        # size of x :[batch_size, sel_len, attention_dim]
        # size of u :[batch_size, attention_dim]
        # uit = tanh(xW+b)
        uit = K.tanh(K.bias_add(K.dot(x, self.W), self.b))
        ait = K.dot(uit, self.u)
        ait = K.squeeze(ait, -1)

        ait = K.exp(ait)

        if mask is not None:
            # Cast the mask to floatX to avoid float64 upcasting in theano
            ait *= K.cast(mask, K.floatx())
        ait /= K.cast(K.sum(ait, axis=1, keepdims=True) + K.epsilon(), K.floatx())
        ait = K.expand_dims(ait)
        weighted_input = x * ait
        output = K.sum(weighted_input, axis=1)

        return output

    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[-1])



tokenizer = Tokenizer(nb_words=MAX_NB_WORDS)
tokenizer.fit_on_texts(flatten_texts)

vocab_size = len(tokenizer.word_index) + 1
embedding_layer = Embedding(vocab_size,
                            EMBEDDING_DIM,
                            #weights=[embedding_matrix],
                            input_length=MAX_PARAGRAPH_LENGTH,
                            trainable=True,
                            mask_zero=True)

paragraph_input = Input(shape=(MAX_PARAGRAPH_LENGTH,), dtype='int32')
embedded_sequences = embedding_layer(paragraph_input)
l_lstm = Bidirectional(GRU(100, return_sequences=True))(embedded_sequences)
l_att = AttLayer(60)(l_lstm)
paragraph_encoder = Model(paragraph_input, l_att)

sequence_input = Input(shape=(MAX_PARAGRAPHS, MAX_PARAGRAPH_LENGTH), dtype='int32')

parag_encoder = TimeDistributed(paragraph_encoder)(sequence_input)
bi_lstm = Bidirectional(LSTM(units=128, return_sequences=True,recurrent_dropout=0.2))(parag_encoder)  # variational biLSTM
final_dense = TimeDistributed(Dense(128, activation="relu"))(bi_lstm)
crf = CRF(n_tags, sparse_target=True)  # CRF layer
out = crf(final_dense)  # output
model = Model(sequence_input, out)
model.compile(optimizer="rmsprop", loss=crf_loss, metrics=[crf_accuracy])

compiling the model leads to following error:

AssertionError: Input mask to CRF must have dim 2 if not None

I am using CRF layer from keras_contrib package

Could you please tell me what I'm doing wrong ?

Solution

This seems to make it run:

final_dense = TimeDistributed(Dense(128, activation="relu"))(bi_lstm)
final_dense = Masking()(final_dense) # insert this
crf = CRF(n_tags, sparse_target=True)  # CRF layer
out = crf(final_dense)

But not entirely sure why.