keras_models.py

def get_learned_scores(**kwargs):
    """
    scores each sentence, then multiply by score before next sequence layer
    """
    sent_len = kwargs.get('sent_len')
    embed_size = kwargs.get('embedding_size')
    seq_len = kwargs.get("seq_len")
    pre_embedded = kwargs.get("pre_embedded", False)
    assert seq_len % sent_len == 0, "sequence length must be a multiple of sentence length"
    sent_per_obs = seq_len // sent_len

    concat_outputs = kwargs.get("concat_outputs", False)

    lstm_units_1 = kwargs.get('lstm_units_1', 16)
    lstm_units_2 = kwargs.get('lstm_cells', 16)

    if pre_embedded:
        inputs = tf.keras.layers.Input(shape=(None, ), name="input")
        embedded = tf.reshape(inputs, (-1, 1200, 768))
    else:
        inputs = tf.keras.layers.Input(shape=(None,), name="input")
        embedded = tf.keras.layers.Embedding(kwargs.get('vocab_size'), embed_size)(inputs)
    reshaped = tf.reshape(embedded, (-1, sent_len, embed_size))
    lstm_level1 = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(lstm_units_1))(reshaped)

    x = tf.keras.layers.Dense(1, activation=None)(lstm_level1)
    logits = tf.reshape(x, (-1, sent_per_obs))
    score = tf.keras.layers.Softmax(name="score")(logits)
    weighted = tf.multiply(lstm_level1, tf.reshape(score, (-1, 1)))

    reshaped_level2 = tf.reshape(weighted, (-1, sent_per_obs, lstm_units_1*2))
    w_average = tf.keras.layers.GlobalAveragePooling1D()(reshaped_level2)

    lstm_level2 = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(lstm_units_2))(reshaped_level2)

    if concat_outputs:
        classifier = tf.keras.layers.Dense(1)(lstm_level2)
        classifier2 = tf.keras.layers.Dense(1)(w_average)
        outputs = tf.keras.layers.concatenate([classifier, classifier2], name="output")
    else:
        classifier = tf.keras.layers.Dense(1, name="output")(lstm_level2)
        classifier2 = tf.keras.layers.Dense(1, name="output_2")(w_average)

        outputs = [classifier, classifier2]

    model = tf.keras.Model(inputs=inputs, outputs=outputs)
    return model