model_architecture.py

    feature_spec = tf_transform_output.transformed_feature_spec() 
    feature_spec.pop(_LABEL_KEY)

    inputs = {key: tf.keras.layers.Input(shape=(max_seq_length), 
                                         name=key, dtype=tf.int32)
              for key in feature_spec.keys()}
  
    input_word_ids = tf.cast(inputs["input_word_ids"], dtype=tf.int32)
    input_mask = tf.cast(inputs["input_mask"], dtype=tf.int32)
    input_type_ids = tf.cast(inputs["input_type_ids"], dtype=tf.int32)
 
    bert_layer = load_bert_layer()
    pooled_output, _ = bert_layer(
        [input_word_ids, 
         input_mask, 
         input_type_ids
        ]
    )
    
    x = tf.keras.layers.Dense(256, activation='relu')(pooled_output)
    dense = tf.keras.layers.Dense(64, activation='relu')(x)
    pred = tf.keras.layers.Dense(1, activation='sigmoid')(dense)
 
    model = tf.keras.Model(
        inputs=[inputs['input_word_ids'], 
                inputs['input_mask'], 
                inputs['input_type_ids']], 
        outputs=pred
    )
    model.compile(loss='binary_crossentropy', 
                        optimizer='adam', 
                        metrics=['accuracy'])