Empty prediction with keras Seq2Seq with attention mechanism

Question

I have a simply seq2seq model with attention mechanism in keras. My problem is that the inference model only gives me empty prediction. However, if I remove the attention it suvessfully gives me the prediction.

This is my model:

# Encoder
encoder_inputs = Input(shape=(max_text_len, ))

# Embedding layer
enc_emb = Embedding(x_voc, embedding_dim,
                    trainable=True)(encoder_inputs)

# Encoder LSTM 1
encoder_lstm1 = Bidirectional(LSTM(latent_dim, return_sequences=True,
                     return_state=True, dropout=0.4,
                     recurrent_dropout=0.4))
(encoder_output1, forward_h1, forward_c1, backward_h1, backward_c1) = encoder_lstm1(enc_emb)

# Encoder LSTM 2
encoder_lstm2 = Bidirectional(LSTM(latent_dim, return_sequences=True,
                     return_state=True, dropout=0.4,
                     recurrent_dropout=0.4))
(encoder_output2, forward_h2, forward_c2, backward_h2, backward_c2) = encoder_lstm2(encoder_output1)

# Encoder LSTM 3
encoder_lstm3 = Bidirectional(LSTM(latent_dim, return_state=True,
                     return_sequences=True, dropout=0.4,
                     recurrent_dropout=0.4))
(encoder_outputs, forward_h, forward_c, backward_h, backward_c) = encoder_lstm3(encoder_output2)

state_h = Concatenate()([forward_h, backward_h])
state_c = Concatenate()([forward_c, backward_c])

# Set up the decoder, using encoder_states as the initial state
decoder_inputs = Input(shape=(None, ))

# Embedding layer
dec_emb_layer = Embedding(y_voc, embedding_dim, trainable=True)
dec_emb = dec_emb_layer(decoder_inputs)


# Decoder LSTM
decoder_lstm = LSTM(latent_dim*2, return_sequences=True,
                    return_state=True, dropout=0.4,
                    recurrent_dropout=0.2)
(decoder_outputs, decoder_fwd_state, decoder_back_state) = \
    decoder_lstm(dec_emb, initial_state=[state_h, state_c])


attention=SeqSelfAttention(attention_activation='sigmoid')(decoder_outputs) 

# Dense layer
decoder_dense = TimeDistributed(Dense(y_voc, activation='softmax'))(attention)

# Define the model
model = Model([encoder_inputs, decoder_inputs], decoder_dense)

This is the inference model.

model = load_model("trained model/model.h5")
encoder_inputs = model.input[0]  # input_1

encoder_outputs, forward_h, forward_c, backward_h, backward_c = model.layers[5].output #Bi-lstm2

state_h_enc = Concatenate()([forward_h, backward_h])
state_c_enc = Concatenate()([forward_c, backward_c])

encoder_states = [state_h_enc, state_c_enc]
encoder_model = Model(encoder_inputs, encoder_states)

decoder_inputs = model.input[1]  # input_2
decoder_state_input_h = Input(shape=(latent_dim*2,), name="input_3")
decoder_state_input_c = Input(shape=(latent_dim*2,), name="input_4")
decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]
decoder_emdedding = model.layers[6](decoder_inputs)
decoder_lstm = model.layers[9]
decoder_outputs, state_h_dec, state_c_dec = decoder_lstm(decoder_emdedding, initial_state=decoder_states_inputs)
decoder_states = [state_h_dec, state_c_dec]

attention = model.layers[-2](decoder_outputs)

decoder_dense = model.layers[-1](attention)
#outputs = decoder_dense(attention)
decoder_model = Model(
    [decoder_inputs] + decoder_states_inputs, [decoder_dense] + decoder_states
)

If I look into I can see that the "problem" lies in decode_sequence()

def decode_sequence(input_seq):

    # Encode the input as state vectors.
    states_value = encoder_model.predict(input_seq)

    # Generate empty target sequence of length 1
    target_seq = np.zeros((1, 1))

    # Populate the first word of target sequence with the start word.
    target_seq[0, 0] = target_word_index['sostok']

    stop_condition = False
    decoded_sentence = ''

    while not stop_condition:
        (output_tokens, h, c) = decoder_model.predict([target_seq]
                + states_value)

        # Sample a token
        sampled_token_index = np.argmax(output_tokens[0, -1, :])
        sampled_token = reverse_target_word_index[sampled_token_index]

        if sampled_token != 'eostok':
            decoded_sentence += ' ' + sampled_token

        # Exit condition: either hit max length or find the stop word.
        if sampled_token == 'eostok' or len(decoded_sentence.split()) \
            >= max_summary_len - 1:
            stop_condition = True

        # Update the target sequence (of length 1)
        target_seq = np.zeros((1, 1))
        target_seq[0, 0] = sampled_token_index

        # Update internal states
        states_value = [h, c]

    return decoded_sentence

When I get the token index with sampled_token_index = np.argmax(output_tokens[0, -1, :]). The argmax is the index for the word "eostok", which is a end of sentence token.

When I try to get the index of row with second-largest value with sampled_token_index = np.argsort(output_tokens[0, -1, :], axis=0)[-2], I get the prediction, albeit very bad one.

Also, instead of using decoder sequence, if I use model.predict() (even if it is wrong for seq2seq model) I get very good prediction. It does not make sense to me.

Do you maybe have an idea why I am having this issues? Is there something wrong with this attention layer:

attention=SeqSelfAttention(attention_activation='sigmoid')(decoder_outputs)

I am using this keras-self-attention

Thanks a lot!