import tensorflow as tf

import bert_layer

def build_model(bert_path, max_seq_length): 
    in_id = tf.keras.layers.Input(shape=(max_seq_length,), name="input_ids")
    in_mask = tf.keras.layers.Input(shape=(max_seq_length,), name="input_masks")
    in_segment = tf.keras.layers.Input(shape=(max_seq_length,), name="segment_ids")

    bert_inputs = [in_id, in_mask, in_segment]
    
    bert_output = bert_layer.BertLayer(bert_path, n_fine_tune_layers=3)(bert_inputs)
    dense = tf.keras.layers.Dense(256, activation='relu')(bert_output)
    pred = tf.keras.layers.Dense(1, activation='sigmoid')(dense)
    
    model = tf.keras.models.Model(
        inputs=bert_inputs, 
        outputs=pred
    )

    model.compile(
        loss='binary_crossentropy',
        optimizer='adam', 
        metrics=['accuracy']
    )
    
    return model