# Copyright 2024 The TensorFlow Ranking Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""Provides an `EstimatorBuilder` for creating a TF-Ranking model estimator.

This class contains the boilerplate that is required to create an estimator for

a TF-Ranking model. The goal is to reduce replicated setups (e.g., transform

function, scoring function) for adopting TF-Ranking. Advanced users can also

derive from this class and further tailor for their needs.

"""

import tensorflow as tf

from tensorflow import estimator as tf_estimator

from tensorflow.python.util import function_utils

from tensorflow_ranking.python import feature

from tensorflow_ranking.python import head

from tensorflow_ranking.python import losses

from tensorflow_ranking.python import metrics

from tensorflow_ranking.python import model

from tensorflow_ranking.python import utils

_METRIC_WEIGHT = "metric_weights_feature_name"

_LOSS_WEIGHT = "loss_weights_feature_name"

# Postfix for names of subscore tensors in GAM.

_SUBSCORE_POSTFIX = "subscore"

# Postfix for names of subweight tensors in GAM.

_SUBWEIGHT_POSTFIX = "subweight"

def _validate_hparams(hparams_in_dict, required_keys, allowed_keys=None):

"""Validate keys in `hparams_in_dict`.

It asserts all of the `required_keys` are presented in `hparams_in_dict`. It

also prints all the allowed keys as info. The main purpose is to allow for

easy debugging, because a mis-spelled allowed key uses the default value

(e.g., None) silently.

A hparams can contain extra keys other than `required_keys` and `allowed_keys`

and this function does nothing about them.

Args:

hparams_in_dict: (dict) A dict with the key in string and value in any type.

required_keys: (list) A list of strings for required keys.

allowed_keys: (list) A list of strings for allowed but not requried keys.

Raises:

ValueError: If any of `required_keys` does not present in `hparams_in_dict`.

"""

required_keys = required_keys or []

allowed_keys = allowed_keys or []

for required_key in required_keys:

if required_key not in hparams_in_dict:

raise ValueError("Required key is missing: '{}'".format(required_key))

for allowed_key in allowed_keys:

if allowed_key in hparams_in_dict:

tf.compat.v1.logging.info("Allowed key is set: {}={}".format(

allowed_key, hparams_in_dict.get(allowed_key)))

else:

tf.compat.v1.logging.info(

"Allowed key is not set: {}".format(allowed_key))

def _validate_function_args(function, required_args):

"""Asserts all of the `required_args` are presented in `function` args.

Args:

function: (function) A python function.

required_args: (list) A list of strings indicating the required args.

Raises:

ValueError: If any of `required_args` does not present in the `function`.

"""

fn_args = function_utils.fn_args(function)

if set(fn_args) != set(required_args):

raise ValueError(

"Function `%s` needs to have the following arguments: %s."

" What were provided are the following: %s." %

(function.__name__, sorted(required_args), sorted(fn_args)))

def _get_metric_pair(key, weight=None, topn=None):

"""Helper function to construct metric name and function."""

name = "".join([

"metric/",

"weighted_" if weight else "",

key,

"_%s" % topn if topn else "",

])

return name, metrics.make_ranking_metric_fn(

key, weights_feature_name=weight, topn=topn)

class EstimatorBuilder(object):

"""Builds a tf.estimator.Estimator for a TF-Ranking model.

An example use case is provided below:

```python

import tensorflow as tf

import tensorflow_ranking as tfr

def scoring_function(context_features, example_features, mode):

# ...

# scoring logic

# ...

return scores # tensors with shape [batch_size, ...]

context_feature_columns = {

"c1": tf.feature_column.numeric_column("c1", shape=(1,))

}

example_feature_columns = {

"e1": tf.feature_column.numeric_column("e1", shape=(1,))

}

hparams = dict(

checkpoint_secs=120,

listwise_inference=False,

loss="softmax_loss",

model_dir="/path/to/your/model_dir/",

num_checkpoints=100)

ranking_estimator = tfr.estimator.EstimatorBuilder(

context_feature_columns,

example_feature_columns,

scoring_function=scoring_function,

hparams=hparams).make_estimator()

```

If you want to customize certain `EstimatorBuilder` behaviors, please create

a subclass of `EstimatorBuilder`, and overwrite related functions. Right now,

we recommend only overwriting the `_eval_metric_fns` for your eval metrics.

For instance, if you need MAP (Mean Average Precision) as your evaluation

metric, you can do the following:

```python

class MyEstimatorBuilder(tfr.estimator.EstimatorBuilder):

def _eval_metric_fns(self):

metric_fns = {}

metric_fns.update({

"metric/ndcg@%d" % topn: tfr.metrics.make_ranking_metric_fn(

tfr.metrics.RankingMetricKey.MAP, topn=topn) for topn in [5, 10]

})

return metric_fns

# Then, you can define your estimator with:

ranking_estimator = MyEstimatorBuilder(

context_feature_columns,

example_feature_columns,

scoring_function=scoring_function,

hparams=hparams).make_estimator()

```

If you really need to overwrite other functions, particularly `_transform_fn`,

`_group_score_fn` and `model_fn`, please be careful because the passed-in

parameters might no longer be used.

"""

def __init__(self,

context_feature_columns,

example_feature_columns,

scoring_function,

transform_function=None,

optimizer=None,

loss_reduction=None,

hparams=None):

"""Constructor.

Args:

context_feature_columns: (dict) Context (aka, query) feature columns.

example_feature_columns: (dict) Example (aka, document) feature columns.

scoring_function: (function) A user-provided scoring function with the

below signatures:

* Args:

`context_features`: (dict) A dict of Tensors with the shape

[batch_size, ...].

`example_features`: (dict) A dict of Tensors with the shape

[batch_size, ...].

`mode`: (`estimator.ModeKeys`) Specifies if this is for training,

evaluation or inference. See ModeKeys.

* Returns: The computed logits, a Tensor of shape [batch_size, 1].

transform_function: (function) A user-provided function that transforms

raw features into dense Tensors with the following signature:

* Args:

`features`: (dict) A dict of Tensors or SparseTensors containing the

raw features from an `input_fn`.

`mode`: (`estimator.ModeKeys`) Specifies if this is for training,

evaluation or inference. See ModeKeys.

* Returns:

`context_features`: (dict) A dict of Tensors with the shape

[batch_size, ...].

`example_features`: (dict) A dict of Tensors with the shape

[batch_size, list_size, ...].

optimizer: (`tf.Optimizer`) An `Optimizer` object for model optimzation.

loss_reduction: (str) An enum of strings indicating the loss reduction

type. See type definition in the `tf.compat.v1.losses.Reduction`.

hparams: (dict) A dict containing model hyperparameters.

Raises:

ValueError: If the `example_feature_columns` is None.

ValueError: If the `scoring_function` is None..

ValueError: If both the `optimizer` and the `hparams["learning_rate"]`

are not specified.

"""

if example_feature_columns is None:

raise ValueError("The `example_feature_columns` is not specified!")

if scoring_function is None:

raise ValueError("The `scoring_function` needs to be specified!")

self._context_feature_columns = context_feature_columns

self._example_feature_columns = example_feature_columns

self._scoring_function = scoring_function

self._transform_function = transform_function

self._hparams = hparams

self._validate_function_args_and_hparams()

if not optimizer and not hparams.get("learning_rate"):

raise ValueError("Please specify either the `optimizer` or the "

"`learning_rate` in `hparams`!")

if optimizer and hparams.get("learning_rate"):

tf.compat.v1.logging.warning("`learning_rate` from `hparams` is ignored "

"as the `optimizer` has been specified!")

self._optimizer = (

optimizer or tf.compat.v1.train.AdagradOptimizer(

learning_rate=hparams.get("learning_rate")))

self._loss_reduction = loss_reduction or tf.compat.v1.losses.Reduction.SUM

def _required_hparam_keys(self):

"""Returns a list of keys for required hparams."""

return [

"checkpoint_secs", "listwise_inference", "loss", "model_dir",

"num_checkpoints"

]

def _allowed_hparam_keys(self):

"""Returns a list of allowed but not required keys for hparams."""

return [_METRIC_WEIGHT, _LOSS_WEIGHT]

def _validate_function_args_and_hparams(self):

"""Validates that the hparams and arguments are all as required."""

_validate_hparams(self._hparams, self._required_hparam_keys(),

self._allowed_hparam_keys())

_validate_function_args(

self._scoring_function,

required_args=["context_features", "example_features", "mode"])

if self._transform_function is not None:

_validate_function_args(

self._transform_function, required_args=["features", "mode"])

def _transform_fn(self, features, mode):

"""Defines the transform fn."""

if self._transform_function is not None:

return self._transform_function(features=features, mode=mode)

if (mode == tf_estimator.ModeKeys.PREDICT and

not self._hparams.get("listwise_inference")):

return feature.encode_pointwise_features(

features=features,

context_feature_columns=self._context_feature_columns,

example_feature_columns=self._example_feature_columns,

mode=mode,

scope="transform_layer")

else:

return feature.encode_listwise_features(

features=features,

context_feature_columns=self._context_feature_columns,

example_feature_columns=self._example_feature_columns,

mode=mode,

scope="transform_layer")

def _get_loss_metric_pair(self, key, weight=None):

"""Constructs metric name and function for a loss."""

name = "".join([

"metric/",

"weighted_" if weight else "",

key,

])

return name, losses.make_loss_metric_fn(key, weights_feature_name=weight)

def _get_loss_metric_pairs(self, key: str, weight=None):

"""Constructs pairs of metric name and function for a loss key."""

if "," in key or ":" in key:

keys_to_weights = utils.parse_keys_and_weights(key)

return [

self._get_loss_metric_pair(key, weight=weight)

for key in keys_to_weights

]

else:

return [self._get_loss_metric_pair(key, weight=weight)]

def _eval_metric_fns(self):

"""Returns a dict from name to metric functions."""

metric_fns = {}

metric_fns.update({

_get_metric_pair(metrics.RankingMetricKey.NDCG, topn=topn)

for topn in [5, 10, None]

})

metric_fns.update({

_get_metric_pair(metrics.RankingMetricKey.MRR, topn=topn)

for topn in [10, None]

})

metric_fns.update({*self._get_loss_metric_pairs(self._hparams.get("loss"))})

if self._hparams.get(_METRIC_WEIGHT):

weight = self._hparams.get(_METRIC_WEIGHT)

tf.compat.v1.logging.info("Metric weight %s=%s" %

(_METRIC_WEIGHT, weight))

metric_fns.update({

_get_metric_pair(

metrics.RankingMetricKey.NDCG, weight=weight, topn=topn)

for topn in [5, 10, None]

})

metric_fns.update({

_get_metric_pair(

metrics.RankingMetricKey.MRR, weight=weight, topn=topn)

for topn in [10, None]

})

metric_fns.update({

*self._get_loss_metric_pairs(

self._hparams.get("loss"), weight=weight)

})

return metric_fns

def _group_score_fn(self, context_features, group_features, mode, params,

config):

"""Returns a groupwise score fn to build `EstimatorSpec`."""

del [params, config] # They are not used.

# Squeeze the group features because we are in univariate mode.

example_features = {}

for k, v in group_features.items():

example_features[k] = tf.squeeze(v, 1)

if self._scoring_function is None:

raise ValueError("The attribute `scoring_function` is being used before"

"being assigned.")

return self._scoring_function(

context_features=context_features,

example_features=example_features,

mode=mode)

def _get_loss_fn(self):

"""Gets the loss function."""

return losses.make_loss_fn(

self._hparams.get("loss"),

weights_feature_name=self._hparams.get(_LOSS_WEIGHT),

reduction=self._loss_reduction)

def _model_fn(self):

"""Returns a model_fn."""

def _train_op_fn(loss):

"""Defines train op used in ranking head."""

update_ops = tf.compat.v1.get_collection(

tf.compat.v1.GraphKeys.UPDATE_OPS)

minimize_op = self._optimizer.minimize(

loss=loss, global_step=tf.compat.v1.train.get_global_step())

train_op = tf.group([update_ops, minimize_op])

return train_op

ranking_head = head.create_ranking_head(

loss_fn=self._get_loss_fn(),

eval_metric_fns=self._eval_metric_fns(),

train_op_fn=_train_op_fn)

return model.make_groupwise_ranking_fn(

group_score_fn=self._group_score_fn,

group_size=1,

transform_fn=self._transform_fn,

ranking_head=ranking_head)

def make_estimator(self):

"""Returns the built `tf.estimator.Estimator` for the TF-Ranking model."""

config = tf_estimator.RunConfig( # pylint: disable=g-deprecated-tf-checker

model_dir=self._hparams.get("model_dir"),

keep_checkpoint_max=self._hparams.get("num_checkpoints"),

save_checkpoints_secs=self._hparams.get("checkpoint_secs"))

return tf_estimator.Estimator(model_fn=self._model_fn(), config=config) # pylint: disable=g-deprecated-tf-checker

def _make_dnn_score_fn(hidden_units,

activation_fn=tf.nn.relu,

dropout=None,

use_batch_norm=False,

batch_norm_moment=0.999):

"""Returns a DNN scoring fn that outputs a score per example.

Args:

hidden_units: (list) Iterable of number hidden units per layer for a DNN

model. All layers are fully connected. Ex. `[64, 32]` means first layer

has 64 nodes and second one has 32.

activation_fn: Activation function applied to each layer. If `None`, will

use `tf.nn.relu`.

dropout: (float) When not `None`, the probability we will drop out a given

coordinate.

use_batch_norm: (bool) If true, use batch normalization after each hidden

layer.

batch_norm_moment: (float) Momentum for the moving average in batch

normalization.

Returns:

A DNN scoring function.

"""

activation_fn = activation_fn or tf.nn.relu

def _scoring_function(context_features, example_features, mode):

"""Defines the DNN-based scoring fn.

Args:

context_features: (dict) A mapping from context feature names to dense 2-D

Tensors of shape [batch_size, ...].

example_features: (dict) A mapping from example feature names to dense 3-D

Tensors of shape [batch_size, list_size, ...].

mode: (`tf.estimator.ModeKeys`) TRAIN, EVAL, or PREDICT.

Returns:

A Tensor of shape [batch_size, 1] containing per-example.

scores.

"""

# Input layer.

with tf.compat.v1.name_scope("input_layer"):

example_input = [

tf.compat.v1.layers.flatten(example_features[name])

for name in sorted(list(example_features.keys()))

]

context_input = [

tf.compat.v1.layers.flatten(context_features[name])

for name in sorted(list(context_features.keys()))

]

# Concat context and example features as input.

input_layer = tf.concat(context_input + example_input, 1)

is_training = (mode == tf_estimator.ModeKeys.TRAIN)

cur_layer = input_layer

# Construct a deep neural network model.

with tf.compat.v1.name_scope("dnn_model"):

if use_batch_norm:

cur_layer = tf.compat.v1.layers.batch_normalization(

cur_layer, training=is_training, momentum=batch_norm_moment)

logits = _feed_forward_network(

cur_layer,

map(int, hidden_units),

output_units=1,

activation_fn=activation_fn,

batch_norm=use_batch_norm,

batch_norm_moment=batch_norm_moment,

dropout=dropout,

is_training=is_training)

tf.compat.v1.summary.scalar("logits_mean",

tf.reduce_mean(input_tensor=logits))

return logits

return _scoring_function

def make_dnn_ranking_estimator(

example_feature_columns,

hidden_units,

context_feature_columns=None,

optimizer=None,

learning_rate=0.05,

loss="approx_ndcg_loss",

loss_reduction=tf.compat.v1.losses.Reduction.SUM_OVER_BATCH_SIZE,

activation_fn=tf.nn.relu,

dropout=None,

use_batch_norm=False,

batch_norm_moment=0.999,

model_dir=None,

checkpoint_secs=120,

num_checkpoints=1000,

listwise_inference=False):

"""Builds an `Estimator` instance with DNN scoring function.

Args:

example_feature_columns: (dict) Example (aka, document) feature columns.

hidden_units: (list) Iterable of number hidden units per layer for a DNN

model. All layers are fully connected. Ex. `[64, 32]` means first layer

has 64 nodes and second one has 32.

context_feature_columns: (dict) Context (aka, query) feature columns.

optimizer: (`tf.Optimizer`) An `Optimizer` object for model optimzation.

learning_rate: (float) Only used if `optimizer` is a string. Defaults to

0.05.

loss: (str) A string to decide the loss function used in training. See

`RankingLossKey` class for possible values.

loss_reduction: (str) An enum of strings indicating the loss reduction type.

See type definition in the `tf.compat.v1.losses.Reduction`.

activation_fn: Activation function applied to each layer. If `None`, will

use `tf.nn.relu`.

dropout: (float) When not `None`, the probability we will drop out a given

coordinate.

use_batch_norm: (bool) Whether to use batch normalization after each hidden

layer.

batch_norm_moment: (float) Momentum for the moving average in batch

normalization.

model_dir: (str) Directory to save model parameters, graph and etc. This can

also be used to load checkpoints from the directory into a estimator to

continue training a previously saved model.

checkpoint_secs: (int) Time interval (in seconds) to save checkpoints.

num_checkpoints: (int) Number of checkpoints to keep.

listwise_inference: (bool) Whether the inference will be performed with the

listwise data format such as `ExampleListWithContext`.

Returns:

An `Estimator` with DNN scoring function.

"""

scoring_function = _make_dnn_score_fn(

hidden_units,

activation_fn=activation_fn,

dropout=dropout,

use_batch_norm=use_batch_norm,

batch_norm_moment=batch_norm_moment)

hparams = dict(

model_dir=model_dir,

learning_rate=learning_rate,

listwise_inference=listwise_inference,

loss=loss,

checkpoint_secs=checkpoint_secs,

num_checkpoints=num_checkpoints)

return EstimatorBuilder(

context_feature_columns,

example_feature_columns,

optimizer=optimizer,

scoring_function=scoring_function,

loss_reduction=loss_reduction,

hparams=hparams).make_estimator()

def _feed_forward_network(x,

hidden_layer_dims,

output_units,

activation_fn=tf.nn.relu,

batch_norm=False,

batch_norm_moment=0.999,

dropout=None,

is_training=None):

"""Defines feed-forward network.

Args:

x: Input tensor.

hidden_layer_dims: Iterable of number hidden units per layer. All layers are

fully connected. Ex. `[64, 32]` means first layer has 64 nodes and second

one has 32.

output_units: (int) Size of output logits from this tower.

activation_fn: Activation function applied to each layer. If `None`, will use

ReLU activation.

batch_norm: Whether to use batch normalization after each hidden layer.

batch_norm_moment: Momentum for the moving average in batch normalization.

dropout: When not `None`, the probability we will drop out a given

coordinate.

is_training: Whether in the training mode.

Returns:

Output tensor.

"""

for layer_width in map(int, hidden_layer_dims):

x = tf.compat.v1.layers.dense(x, units=layer_width)

if batch_norm:

x = tf.compat.v1.layers.batch_normalization(

x, training=is_training, momentum=batch_norm_moment)

x = activation_fn(x)

if dropout:

x = tf.compat.v1.layers.dropout(

inputs=x, rate=dropout, training=is_training)

output = tf.compat.v1.layers.dense(x, units=output_units)

return output

def _make_gam_score_fn(context_hidden_units,

example_hidden_units,

activation_fn=tf.nn.relu,

dropout=None,

batch_norm=False,

batch_norm_moment=0.999):

"""Returns a scoring fn that outputs a score per example."""

activation_fn = activation_fn or tf.nn.relu

def _scoring_fn(context_features, example_features, mode):

"""Defines the scoring fn for GAM.

Args:

context_features: (dict) A mapping from context feature names to dense 2-D

Tensors of shape [batch_size, ...].

example_features: (dict) A mapping from example feature names to dense 3-D

Tensors of shape [batch_size, list_size, ...].

mode: (`tf.estimator.ModeKeys`) TRAIN, EVAL, or PREDICT.

Returns:

A Tensor of shape [batch_size, 1] containing per-example scores.

"""

# Input layer.

example_feature_names = sorted(list(example_features.keys()))

context_feature_names = sorted(list(context_features.keys()))

with tf.compat.v1.name_scope("input_layer"):

example_input = [(name,

tf.compat.v1.layers.flatten(example_features[name]))

for name in sorted(list(example_feature_names))]

context_input = [(name,

tf.compat.v1.layers.flatten(context_features[name]))

for name in sorted(list(context_feature_names))]

is_training = (mode == tf_estimator.ModeKeys.TRAIN)

# Construct a tower for each example feature. Each tower outputs a

# scalar value as the sub-score. All sub-scores are

# [batch_size * list_size, 1]-shaped tensors and are stored in

# `sub_logits_list` as a `feature_num`-sized list.

with tf.compat.v1.name_scope("example_feature_towers"):

sub_logits_list = []

for name, input_layer in example_input:

with tf.compat.v1.name_scope("{}_tower".format(name)):

cur_layer = input_layer

if batch_norm:

cur_layer = tf.compat.v1.layers.batch_normalization(

cur_layer, training=is_training, momentum=batch_norm_moment)

sub_logits = _feed_forward_network(

cur_layer,

map(int, example_hidden_units),

output_units=1,

activation_fn=activation_fn,

batch_norm=batch_norm,

batch_norm_moment=batch_norm_moment,

dropout=dropout,

is_training=is_training)

sub_logits = tf.identity(

sub_logits, name="{}_{}".format(name, _SUBSCORE_POSTFIX))

sub_logits_list.append(sub_logits)

# Construct a tower for each context feature. Each tower outputs a

# weighting vector of `feature_num`-dim where `feature_num` is the number

# of example features. All the vectors are

# [batch_size * list_size, feature_num] tensors and are stored in

# `sub_weights_list` with length of number of context feature.

sub_weights_list = []

if context_input:

# Construct a tower per context features.

with tf.compat.v1.name_scope("context_feature_towers"):

feature_num = len(sub_logits_list)

for name, input_layer in context_input:

with tf.compat.v1.name_scope("{}_tower".format(name)):

cur_layer = input_layer

if batch_norm:

cur_layer = tf.compat.v1.layers.batch_normalization(

cur_layer, training=is_training, momentum=batch_norm_moment)

sub_weights = _feed_forward_network(

cur_layer,

map(int, context_hidden_units),

output_units=feature_num,

activation_fn=activation_fn,

batch_norm=batch_norm,

batch_norm_moment=batch_norm_moment,

dropout=dropout,

is_training=is_training)

sub_weights = tf.math.softmax(

sub_weights, name="{}_{}".format(name, _SUBWEIGHT_POSTFIX))

sub_weights_list.append(sub_weights)

# Construct an additive model from the outputs of all example feature towers

# `sub_logits_list` weighted by outputs of all context feature towers

# `sub_weights_list`. If no context features are provided, the outputs will

# simply be the sum of `sub_logits_list`.

if sub_weights_list:

sub_logits = tf.concat(sub_logits_list, axis=-1)

feature_weights = tf.math.add_n(sub_weights_list)

logits = tf.math.reduce_sum(

input_tensor=sub_logits * feature_weights, axis=-1)

else:

logits = tf.math.add_n(sub_logits_list)

tf.compat.v1.summary.scalar("logits_mean",

tf.reduce_mean(input_tensor=logits))

return logits

return _scoring_fn

# TODO: Attach the link to the paper.

class GAMEstimatorBuilder(EstimatorBuilder):

"""Builds a TFR estimator with subscore signatures of GAM models.

Neural Generalized Additive Ranking Model is an additive ranking model.

See the paper (https://arxiv.org/abs/2005.02553) for more details.

For each example x with n features (x_1, x_2, ..., x_n), the ranking score is:

F(x) = f1(x_1) + f2(x_2) + ... + fn(x_n)

where each feature is scored by a corresponding submodel, and the overall

ranking score is the sum of all the submodels' outputs. Each submodel is a

standalone feed-forward network.

When there are m context features (c_1, c_2, ..., c_m), the ranking score

will be determined by:

F(c, x) = w1(c) * f1(x_1) + w2(c) * f2(x_2) + ... + wn(c) * fn(x_n)

where (w1(c), w2(c), ..., wn(c)) is a weighting vector determined solely by

context features. For each context feature c_j, a feed-forward submodel is

constructed to derive a weighting vector (wj1(c_j), wj2(c_j), ..., wjn(c_j)).

The final weighting vector is the sum of the output of all the context

features' submodels.

The model is implicitly interpretable as the contribution of each feature to

the final ranking score can be easily visualized. However, the model does not

have higher-order inter-feature interactions and hence may not have

performance as good as the fully-connected DNN.

The output of each example feature's submodel can be retrieved by tensor

named `{feature_name}_subscore`. The output of each context feature's submodel

is a n-dimensional vector and can be retrieved by tensor named

`{feature_name}_subweight`.

"""

def _model_fn(self):

"""Wraps model_fn with additional signatures of subscores."""

def _gam_model_fn(features, labels, mode, params, config):

"""Redefines the model_fn for GAM to include subscore signatures."""

estimator_spec = super(GAMEstimatorBuilder,

self)._model_fn()(features, labels, mode, params,

config)

if mode == tf_estimator.ModeKeys.PREDICT:

# Export subscores of each feature. Find nodes ending with

# `_SUBSCORE_POSTFIX` and `_SUBWEIGHT_POSTFIX` and create signatures

# with their corresponding tensors as outputs. Signatures for example

# feature sub-scores are regression signatures, and signatures for

# context feature weighting vectors are prediction signatures.

subscore_signatures = {}

for node in tf.compat.v1.get_default_graph().as_graph_def().node:

if node.name.endswith(_SUBSCORE_POSTFIX):

subscore_name = node.name[node.name.rfind("/") + 1:]

subscore_tensor = (

tf.compat.v1.get_default_graph().get_tensor_by_name(

"{}:0".format(node.name)))

subscore_signatures[subscore_name] = (

tf_estimator.export.RegressionOutput(subscore_tensor)) # pylint: disable=g-deprecated-tf-checker

elif node.name.endswith(_SUBWEIGHT_POSTFIX):

subscore_name = node.name[node.name.rfind("/") + 1:]

subscore_tensor = (

tf.compat.v1.get_default_graph().get_tensor_by_name(

"{}:0".format(node.name)))

subscore_signatures[subscore_name] = (

tf_estimator.export.PredictOutput(subscore_tensor)) # pylint: disable=g-deprecated-tf-checker

estimator_spec.export_outputs.update(subscore_signatures)

return estimator_spec

return _gam_model_fn

def make_gam_ranking_estimator(

example_feature_columns,

example_hidden_units,

context_feature_columns=None,

context_hidden_units=None,

optimizer=None,

learning_rate=0.05,

loss="approx_ndcg_loss",

loss_reduction=tf.compat.v1.losses.Reduction.SUM_OVER_BATCH_SIZE,

activation_fn=tf.nn.relu,

dropout=None,

use_batch_norm=False,

batch_norm_moment=0.999,

model_dir=None,

checkpoint_secs=120,

num_checkpoints=1000,

listwise_inference=False):

"""Builds an `Estimator` instance with GAM scoring function.

See the comment of `GAMEstimatorBuilder` class for more details.

Args:

example_feature_columns: (dict) A dict containing all the example feature

columns used by the model. Keys are feature names, and values are

instances of classes derived from `_FeatureColumn`.

example_hidden_units: (list) Iterable of number hidden units per layer for

example features. All layers are fully connected. Ex. `[64, 32]` means

first layer has 64 nodes and second one has 32.

context_feature_columns: (dict) A dict containing all the context feature

columns used by the model. See `example_feature_columns`.

context_hidden_units: (list) Iterable of number hidden units per layer for

context features. See `example_hidden_units`.

optimizer: (`tf.Optimizer`) An `Optimizer` object for model optimzation. If

`None`, an Adagard optimizer with `learning_rate` will be created.

learning_rate: (float) Only used if `optimizer` is a string. Defaults to

0.05.

loss: (str) A string to decide the loss function used in training. See

`RankingLossKey` class for possible values.

loss_reduction: (str) An enum of strings indicating the loss reduction type.

See type definition in the `tf.compat.v1.losses.Reduction`.

activation_fn: Activation function applied to each layer. If `None`, will

use `tf.nn.relu`.

dropout: (float) When not `None`, the probability we will drop out a given

coordinate.

use_batch_norm: (bool) Whether to use batch normalization after each hidden

layer.

batch_norm_moment: (float) Momentum for the moving average in batch

normalization.

model_dir: (str) Directory to save model parameters, graph and etc. This can

also be used to load checkpoints from the directory into a estimator to

continue training a previously saved model.

checkpoint_secs: (int) Time interval (in seconds) to save checkpoints.

num_checkpoints: (int) Number of checkpoints to keep.

listwise_inference: (bool) Whether the inference will be performed with the

listwise data format such as `ExampleListWithContext`.

Returns:

An `Estimator` with GAM scoring function.

"""

scoring_function = _make_gam_score_fn(

context_hidden_units,

example_hidden_units,

activation_fn=activation_fn,

dropout=dropout,

batch_norm=use_batch_norm,

batch_norm_moment=batch_norm_moment)

hparams = dict(

model_dir=model_dir,

learning_rate=learning_rate,

listwise_inference=listwise_inference,

loss=loss,

checkpoint_secs=checkpoint_secs,

num_checkpoints=num_checkpoints)

return GAMEstimatorBuilder(

context_feature_columns,

example_feature_columns,

optimizer=optimizer,

scoring_function=scoring_function,

loss_reduction=loss_reduction,

hparams=hparams).make_estimator()