import tensorflow as tf
from sklearn.model_selection import train_test_split
from bert4keras.models import build_transformer_model
from bert4keras.tokenizers import Tokenizer
from bert4keras.backend import keras
tf.config.experimental.set_memory_growth(device=tf.config.experimental.list_physical_devices(device_type='GPU')[0], enable=True)
embed_dim = 64
NEG, batch_size = 20, 128
config_path = 'chinese_roberta_L-6_H-384_A-12/bert_config.json'
checkpoint_path = 'chinese_roberta_L-6_H-384_A-12/bert_model.ckpt'
query_bert = build_transformer_model(config_path, checkpoint_path, return_keras_model=False).model
query_layer = keras.layers.Dropout(0.1)(query_bert.output)
query_layer = keras.layers.Dense(128, activation='relu', kernel_regularizer='l2', name="query_tower")(query_layer)
doc_bert = build_transformer_model(config_path, checkpoint_path, return_keras_model=False).model
for layer in doc_bert.layers:
layer.name = layer.name + str("_doc")
doc_layer = keras.layers.Dropout(0.1)(doc_bert.output)
doc_layer = keras.layers.Dense(128, activation='relu', kernel_regularizer='l2', name="doc_tower")(doc_layer)
output = keras.layers.Dot(axes=1)([query_layer, doc_layer])
# output = tf.keras.layers.Dense(1, activation='sigmoid')(output)
output = keras.layers.Dense(2, activation='softmax')(output)
model = keras.models.Model(query_bert.input+doc_bert.input, output)
model.compile(loss="categorical_crossentropy", metrics=['acc' ], optimizer='RMSprop')
# query tower
query_model = keras.Model(inputs=query_bert.input, outputs=query_layer)
# doc tower
doc_model = keras.Model(inputs=doc_bert.input, outputs=doc_layer)
print("[INFO] training model...")
model.fit(
train_generator.forfit(),
steps_per_epoch=len(train_generator),
epochs=2, verbose=1)
