Scalable, neural learning to rank (LTR) models
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_ranking as tfr
# Prep data
ds = tfds.load("mslr_web/10k_fold1", split="train")
ds = ds.map(lambda feature_map: {
"_mask": tf.ones_like(feature_map["label"], dtype=tf.bool),
**feature_map
})
ds = ds.shuffle(buffer_size=1000).padded_batch(batch_size=32)
ds = ds.map(lambda feature_map: (
feature_map, tf.where(feature_map["_mask"], feature_map.pop("label"), -1.)))
# Create a model
inputs = {
"float_features": tf.keras.Input(shape=(None, 136), dtype=tf.float32)
}
norm_inputs = [tf.keras.layers.BatchNormalization()(x) for x in inputs.values()]
x = tf.concat(norm_inputs, axis=-1)
for layer_width in [128, 64, 32]:
x = tf.keras.layers.Dense(units=layer_width)(x)
x = tf.keras.layers.Activation(activation=tf.nn.relu)(x)
scores = tf.squeeze(tf.keras.layers.Dense(units=1)(x), axis=-1)
# Compile and train
model = tf.keras.Model(inputs=inputs, outputs=scores)
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
loss=tfr.keras.losses.SoftmaxLoss(),
metrics=tfr.keras.metrics.get("ndcg", topn=5, name="NDCG@5"))
model.fit(ds, epochs=3)
Run in a Notebook
TensorFlow Ranking is an open-source library for developing scalable, neural learning to rank (LTR) models. Ranking models are typically used in search and recommendation systems, but have also been successfully applied in a wide variety of fields, including machine translation, dialogue systems e-commerce, SAT solvers, smart city planning, and even computational biology.
A ranking model takes a list of items (web pages, documents, products, movies, etc.) and generates a list in an optimized order, such as most relevant items on top and the least relevant items at the bottom, usually in response to a user query:
This library supports standard pointwise, pairwise, and listwise loss functions for LTR models. It also supports a wide range of ranking metrics, including Mean Reciprocal Rank (MRR) and Normalized Discounted Cumulative Gain (NDCG), so you can evaluate and compare these approaches for your ranking task. The Ranking library also provides functions for enhanced ranking approaches that are researched, tested, and built by machine learning engineers at Google.
Get started with the TensorFlow Ranking library by checking out the tutorial. Learn more about the capabilities of the library by reading the Overview Check out the source code for TensorFlow Ranking on GitHub.
