bclassifier.py

input_word_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                       name="input_word_ids")
input_mask = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                   name="input_mask")
segment_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                    name="segment_ids")
bert_inputs = [input_word_ids, input_mask, segment_ids]

pooled_output, _ = bert_layer(bert_inputs) 

dense = tf.keras.layers.Dense(256, activation='relu')(pooled_output)

pred = tf.keras.layers.Dense(1, activation='sigmoid')(dense)

model = Model(inputs=bert_inputs, outputs=pred)