infer.py 809 Bytes
Newer Older
chenpangpang's avatar
chenpangpang committed
1
2
import os
os.environ['TF_USE_LEGACY_KERAS'] = '1'
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text as text
text_input = tf.keras.layers.Input(shape=(), dtype=tf.string)
preprocessor = hub.KerasLayer(
    "https://kaggle.com/models/tensorflow/bert/TensorFlow2/en-uncased-preprocess/3")
encoder_inputs = preprocessor(text_input)
encoder = hub.KerasLayer(
    "https://kaggle.com/models/tensorflow/bert/TensorFlow2/en-uncased-l-12-h-768-a-12/3",
    trainable=True)
outputs = encoder(encoder_inputs)
pooled_output = outputs["pooled_output"]      # [batch_size, 768].
sequence_output = outputs["sequence_output"]  # [batch_size, seq_length, 768].
embedding_model = tf.keras.Model(text_input, pooled_output)
sentences = tf.constant(["(your text here)"])
print(embedding_model(sentences))