from tensorflow.python.keras.layers import Dense, Input, Dropout
from tensorflow.python.keras.optimizers import Adam
from tensorflow.python.keras.models import Model
from sklearn.metrics import roc_auc_score

def build_model(transformer_layer, loss="binary_crossentropy", max_len=None):
    input_ids = Input(shape=(max_len,), dtype=tf.int32, name="input_ids")
    attention_mask = Input(shape=(max_len,), dtype=tf.int32, name="attention_mask")

    sequence_output = transformer_layer([input_ids, attention_mask])
    hidden_state = sequence_output["last_hidden_state"]
    cls_token = hidden_state[:, 0, :]  # cls_token은 첫 번째

    x = Dropout(0.35)(cls_token)
    out = Dense(1, activation="sigmoid")(x)

    model = Model(inputs=[input_ids, attention_mask], outputs=out)
    model.compile(
             Adam(learning_rate=3e-5), loss=loss, metrics=[tf.keras.metrics.AUC()]
    )

    return model