BERT Listwise Ranking - tensorflow

Question

I am trying to make a listwise ranking model similar to the Tensorflow Listwise Ranking Tutorial, however, the only difference is changing how the embedding is done. I want to use the BERT embedding model shown in Tensorflow Hub.
Here is my code:
import tensorflow_hub as hub
import tensorflow_text as text
import tensorflow as tf
import tensorflow_datasets as tfds
import numpy as np
import tensorflow_ranking as tfr
import tensorflow_recommenders as tfrs
preprocessor = hub.KerasLayer("https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3")
encoder = hub.KerasLayer("https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/4",trainable=True)
text_input = tf.keras.layers.Input(shape=(), dtype=tf.string)
encoder_inputs = preprocessor(text_input)
outputs = encoder(encoder_inputs)
pooled_output = outputs["pooled_output"]
embedding_model1 = tf.keras.Model({"user_id":text_input}, pooled_output)
text_input = tf.keras.layers.Input(shape=(), dtype=tf.string)
encoder_inputs = preprocessor(text_input)
outputs = encoder(encoder_inputs)
pooled_output = outputs["pooled_output"]
embedding_model2 = tf.keras.Model({"movie_title": text_input}, pooled_output)
ratings = tfds.load("movielens/100k-ratings", split="train")
movies = tfds.load("movielens/100k-movies", split="train")
ratings = ratings.map(lambda x: {
"movie_title": x["movie_title"],
"user_id": x["user_id"],
"user_rating": x["user_rating"],
})
movies = movies.map(lambda x: x["movie_title"])
unique_movie_titles = np.unique(np.concatenate(list(movies.batch(1000))))
unique_user_ids = np.unique(np.concatenate(list(ratings.batch(1_000).map(
lambda x: x["user_id"]))))
tf.random.set_seed(42)
# Split between train and tests sets, as before.
shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)
train = shuffled.take(80_000)
test = shuffled.skip(80_000).take(20_000)
# We sample 50 lists for each user for the training data. For each list we
# sample 5 movies from the movies the user rated.
train = tfrs.examples.movielens.sample_listwise(
train,
num_list_per_user=50,
num_examples_per_list=5,
seed=42
)
test = tfrs.examples.movielens.sample_listwise(
test,
num_list_per_user=1,
num_examples_per_list=5,
seed=42
)
epochs = 10
cached_train = train.shuffle(100_000).batch(32).cache()
cached_test = test.batch(16).cache()
class RankingModel(tfrs.Model):
def __init__(self, loss):
super().__init__()
# Compute embeddings for users.
self.user_embeddings = embedding_model1
# Compute embeddings for movies.
self.movie_embeddings = embedding_model2
# Compute predictions.
self.score_model = tf.keras.Sequential([
# Learn multiple dense layers.
tf.keras.layers.Dense(256, activation="relu"),
tf.keras.layers.Dense(64, activation="relu"),
# Make rating predictions in the final layer.
tf.keras.layers.Dense(1)
])
self.task = tfrs.tasks.Ranking(
loss=loss,
metrics=[
tfr.keras.metrics.NDCGMetric(name="ndcg_metric"),
tf.keras.metrics.RootMeanSquaredError()
]
)
def call(self, features):
# We first convert the id features into embeddings.
# User embeddings are a [batch_size, embedding_dim] tensor.
user_embeddings = self.user_embeddings(features["user_id"])
# Movie embeddings are a [batch_size, num_movies_in_list, embedding_dim]
# tensor.
movie_embeddings = self.movie_embeddings(features["movie_title"])
# We want to concatenate user embeddings with movie emebeddings to pass
# them into the ranking model. To do so, we need to reshape the user
# embeddings to match the shape of movie embeddings.
list_length = features["movie_title"].shape[1]
user_embedding_repeated = tf.repeat(
tf.expand_dims(user_embeddings, 1), [list_length], axis=1)
# Once reshaped, we concatenate and pass into the dense layers to generate
# predictions.
concatenated_embeddings = tf.concat(
[user_embedding_repeated, movie_embeddings], 2)
return self.score_model(concatenated_embeddings)
def compute_loss(self, features, training=False):
labels = features.pop("user_rating")
scores = self(features)
return self.task(
labels=labels,
predictions=tf.squeeze(scores, axis=-1),
)
mse_model = RankingModel(tf.keras.losses.MeanSquaredError())
mse_model.compile(optimizer=tf.keras.optimizers.Adagrad(0.1))
mse_model.fit(cached_train, epochs=1, verbose=False)
mse_model.summary()
I think the issue might be the fact that BERT from Tensorflow Hub can only take in 1D tensors like (N,). How can I change the code so I can make a listwise ranking model that uses BERT as the embedder? Here is the error message I get:
WARNING:tensorflow:Model was constructed with shape (None,) for input KerasTensor(type_spec=TensorSpec(shape=(None,), dtype=tf.string, name='input_39'), name='input_39', description="created by layer 'input_39'"), but it was called on an input with incompatible shape (None, 5).
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-117-f6e1e67627f1> in <module>
127 mse_model = RankingModel(tf.keras.losses.MeanSquaredError())
128 mse_model.compile(optimizer=tf.keras.optimizers.Adagrad(0.1))
--> 129 mse_model.fit(cached_train, epochs=1, verbose=False)
130 mse_model.summary()
7 frames
/usr/local/lib/python3.7/dist-packages/tensorflow_hub/keras_layer.py in <lambda>()
70 training = False
71 ag__.if_stmt(ag__.ld(self).trainable, if_body_2, else_body_2, get_state_2, set_state_2, ('training',), 1)
---> 72 result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)
73 result = ag__.Undefined('result')
74 ag__.if_stmt(ag__.not_(ag__.ld(self)._has_training_argument), if_body_3, else_body_3, get_state_3, set_state_3, ('result', 'training'), 1)
ValueError: in user code:
File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1160, in train_function *
return step_function(self, iterator)
File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1146, in step_function **
outputs = model.distribute_strategy.run(run_step, args=(data,))
File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1135, in run_step **
outputs = model.train_step(data)
File "/usr/local/lib/python3.7/dist-packages/tensorflow_recommenders/models/base.py", line 68, in train_step
loss = self.compute_loss(inputs, training=True)
File "<ipython-input-117-f6e1e67627f1>", line 120, in compute_loss
scores = self(features)
File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/tmp/__autograph_generated_filedc6sn1ck.py", line 11, in tf__call
movie_embeddings = ag__.converted_call(ag__.ld(self).movie_embeddings, (ag__.ld(features)['movie_title'],), None, fscope)
File "/tmp/__autograph_generated_file6e8815fm.py", line 74, in tf__call
ag__.if_stmt(ag__.not_(ag__.ld(self)._has_training_argument), if_body_3, else_body_3, get_state_3, set_state_3, ('result', 'training'), 1)
File "/tmp/__autograph_generated_file6e8815fm.py", line 72, in else_body_3
result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)
File "/tmp/__autograph_generated_file6e8815fm.py", line 72, in <lambda>
result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)
ValueError: Exception encountered when calling layer "ranking_model_13" " f"(type RankingModel).
in user code:
File "<ipython-input-117-f6e1e67627f1>", line 101, in call *
movie_embeddings = self.movie_embeddings(features["movie_title"])
File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler **
raise e.with_traceback(filtered_tb) from None
File "/tmp/__autograph_generated_file6e8815fm.py", line 74, in tf__call
ag__.if_stmt(ag__.not_(ag__.ld(self)._has_training_argument), if_body_3, else_body_3, get_state_3, set_state_3, ('result', 'training'), 1)
File "/tmp/__autograph_generated_file6e8815fm.py", line 72, in else_body_3
result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)
File "/tmp/__autograph_generated_file6e8815fm.py", line 72, in <lambda>
result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)
ValueError: Exception encountered when calling layer "keras_layer_8" " f"(type KerasLayer).
in user code:
File "/usr/local/lib/python3.7/dist-packages/tensorflow_hub/keras_layer.py", line 237, in call *
result = smart_cond.smart_cond(training,
ValueError: Could not find matching concrete function to call loaded from the SavedModel. Got:
Positional arguments (3 total):
* <tf.Tensor 'inputs:0' shape=(None, 5) dtype=string>
* False
* None
Keyword arguments: {}
Expected these arguments to match one of the following 4 option(s):
Option 1:
Positional arguments (3 total):
* TensorSpec(shape=(None,), dtype=tf.string, name='sentences')
* False
* None
Keyword arguments: {}
Option 2:
Positional arguments (3 total):
* TensorSpec(shape=(None,), dtype=tf.string, name='sentences')
* True
* None
Keyword arguments: {}
Option 3:
Positional arguments (3 total):
* TensorSpec(shape=(None,), dtype=tf.string, name='inputs')
* False
* None
Keyword arguments: {}
Option 4:
Positional arguments (3 total):
* TensorSpec(shape=(None,), dtype=tf.string, name='inputs')
* True
* None
Keyword arguments: {}
Call arguments received by layer "keras_layer_8" " f"(type KerasLayer):
• inputs=tf.Tensor(shape=(None, 5), dtype=string)
• training=False
Call arguments received by layer "ranking_model_13" " f"(type RankingModel):
• features={'user_id': 'tf.Tensor(shape=(None,), dtype=string)', 'movie_title': 'tf.Tensor(shape=(None, 5), dtype=string)'}