core-power / 2021ccf-iqiyi Goto Github PK

def build_model():
    token_ids = Input(shape=(None,))  
    segment_ids = Input(shape=(None,))

    bert = build_transformer_model(
        config_path=config_path,
        checkpoint_path=checkpoint_path,
        model='bert',
        return_keras_model=False
    )

    output = Lambda(lambda x: x[:, 0])(bert.model([token_ids, segment_ids]))  # 取出output中的第一个值，代表[cls]的值

    output_love = Dense(units=emotion_classes,activation='softmax',kernel_initializer=bert.initializer,name='love_classes')(output)  # (?,4)
    output_happy = Dense(units=emotion_classes,activation='softmax',kernel_initializer=bert.initializer,name='happy_classes')(output)  # (?,4)
    output_scare = Dense(units=emotion_classes,activation='softmax',kernel_initializer=bert.initializer,name='scare_classes')(output)  # (?,4)
    output_angry = Dense(units=emotion_classes,activation='softmax',kernel_initializer=bert.initializer,name='angry_classes')(output)  # (?,4)
    output_afraid = Dense(units=emotion_classes,activation='softmax',kernel_initializer=bert.initializer,name='afraid_classes')(output)  # (?,4)
    output_sad = Dense(units=emotion_classes,activation='softmax',kernel_initializer=bert.initializer,name='sad_classes')(output)  # (?,4)


    model = keras.models.Model([token_ids, segment_ids], [output_love, output_happy,output_scare,output_angry,output_afraid,output_sad])
    model.summary()

    model.compile(
        loss='categorical_crossentropy',
        optimizer=Adam(lr),
        metrics=['categorical_accuracy']
    )
    return model