Sentence Transformer in tf-transformers

  • This is a simple tutorial to demonstrate how SentenceTransformer models has been integrated to tf-transformers and how to use it

  • The following tutorial is applicable to all supported SentenceTransformer models.

Load Sentence-t5 model

import tensorflow as tf
from tf_transformers.models import SentenceTransformer

model_name = 'sentence-transformers/sentence-t5-base' # Load any sentencetransformer model here
model = SentenceTransformer.from_pretrained(model_name)

Metal device set to: Apple M1

2022-04-02 17:00:46.902688: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:305] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.
2022-04-02 17:00:46.902786: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:271] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)
INFO:absl:Successful ✅✅: Model checkpoints matched and loaded from /Users/sarathrnair/.cache/huggingface/hub/tftransformers__sentence-t5-base-sentence-transformers.main.d64dbdc4c8c15637da4215b81f38af99d48a586c/ckpt-1
INFO:absl:Successful ✅: Loaded model from tftransformers/sentence-t5-base-sentence-transformers

Whats my model input?

  • All models in tf-transformers are designed with full connections. All you need is model.input if its a LegacyModel/tf.keras.Model or model.model_inputs if its a LegacyLayer/tf.keras.layers.Layer

model.input

{'input_ids': <KerasTensor: shape=(None, None) dtype=int32 (created by layer 'input_ids')>,
 'input_mask': <KerasTensor: shape=(None, None) dtype=int32 (created by layer 'input_mask')>}

Whats my model output?

  • All models in tf-transformers are designed with full connections. All you need is model.output if its a LegacyModel/tf.keras.Model or model.model_outputs if its a LegacyLayer/tf.keras.layers.Layer

model.output

{'sentence_vector': <KerasTensor: shape=(None, 768) dtype=float32 (created by layer 'tf_transformers/t5_encoder')>}

Sentence vectors

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_name)

text = ['This is a sentence to get vector', 'This one too']
inputs = tokenizer(text, return_tensors='tf', padding=True)

inputs_tf = {'input_ids': inputs['input_ids'], 'input_mask': inputs['attention_mask']}
outputs_tf = model(inputs_tf)
print("Sentence vector", outputs_tf['sentence_vector'].shape)

Sentence vector (2, 768)

Serialize as usual and load it

  • Serialize, load and assert outputs with non serialized (tf.keras.Model)

model_dir = 'MODELS/sentence_t5'
model.save_transformers_serialized(model_dir)

loaded = tf.saved_model.load(model_dir)
model = loaded.signatures['serving_default']

outputs_tf_serialized = model(**inputs_tf)

tf.debugging.assert_near(outputs_tf['sentence_vector'], outputs_tf_serialized['sentence_vector'])

2022-04-02 17:09:27.002497: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
WARNING:absl:Found untraced functions such as tf_transformers/t5_encoder_layer_call_fn, tf_transformers/t5_encoder_layer_call_and_return_conditional_losses, grt5_dense_layer_layer_call_fn, grt5_dense_layer_layer_call_and_return_conditional_losses, dropout_2_layer_call_fn while saving (showing 5 of 880). These functions will not be directly callable after loading.

INFO:tensorflow:Assets written to: MODELS/sentence_t5/assets

INFO:tensorflow:Assets written to: MODELS/sentence_t5/assets
2022-04-02 17:09:32.964144: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz
2022-04-02 17:09:32.965446: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:112] Plugin optimizer for device_type GPU is enabled.