From 272154f951c9ce1fe0bea0bdf12cb71ac64ce1ca Mon Sep 17 00:00:00 2001 From: luozhouyang Date: Sun, 24 Jan 2021 21:32:46 +0800 Subject: [PATCH] Update README --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 4676458..e30464d 100644 --- a/README.md +++ b/README.md @@ -49,7 +49,7 @@ def build_bert_classify_model(pretrained_model_dir, trainable=True, **kwargs): bert = Bert.from_pretrained(pretrained_model_dir, **kwargs) bert.trainable = trainable - _, pooled_output, _, _ = bert(inputs=(input_ids, segment_ids)) + sequence_outputs, pooled_output = bert(inputs=(input_ids, segment_ids)) outputs = tf.keras.layers.Dense(2, name='output')(pooled_output) model = tf.keras.Model(inputs=[input_ids, segment_ids], outputs=outputs) model.compile(loss='binary_cross_entropy', optimizer='adam') @@ -87,7 +87,7 @@ def build_albert_classify_model(pretrained_model_dir, trainable=True, **kwargs): albert = Albert.from_pretrained(pretrained_model_dir, **kwargs) albert.trainable = trainable - _, pooled_output, _, _ = albert(inputs=(input_ids, segment_ids)) + sequence_outputs, pooled_output = albert(inputs=(input_ids, segment_ids)) outputs = tf.keras.layers.Dense(2, name='output')(pooled_output) model = tf.keras.Model(inputs=[input_ids, segment_ids], outputs=outputs) model.compile(loss='binary_cross_entropy', optimizer='adam')