def test_create_pipeline_with_misspecified_args(self):
hparams = _make_hparams(
train_input_pattern=self._data_file,
eval_input_pattern=self._data_file,
model_dir=self._model_dir)
# The `estimator` cannot be None.
with self.assertRaises(ValueError):
_ = pipeline.RankingPipeline(
_context_feature_columns(),
_example_feature_columns(),
hparams=hparams,
estimator=None,
label_feature_name="utility",
label_feature_type=tf.float32)
# A dict is not an `Estimator`.
estimator = {}
with self.assertRaises(ValueError):
_ = pipeline.RankingPipeline(
_context_feature_columns(),
_example_feature_columns(),
hparams=hparams,
estimator=estimator,
label_feature_name="utility",
label_feature_type=tf.float32)
def test_train_and_eval(self):
def _trainable_score_fn(context_features, example_features, mode):
del context_features, mode
input_layer = tf.ones_like(example_features["f1"])
return tf.compat.v1.layers.dense(input_layer, units=1)
hparams = _make_hparams(
train_input_pattern=self._data_file,
eval_input_pattern=self._data_file,
model_dir=self._model_dir,
list_size=None,
num_train_steps=1,
num_eval_steps=1)
estimator = tfr.estimator.EstimatorBuilder(
_context_feature_columns(),
_example_feature_columns(),
_trainable_score_fn,
hparams=hparams).make_estimator()
pip = pipeline.RankingPipeline(
_context_feature_columns(),
_example_feature_columns(),
hparams=hparams,
estimator=estimator,
label_feature_name="utility",
label_feature_type=tf.float32)
pip.train_and_eval()
hparams["listwise_inference"] = True
pip_listwise_export = pipeline.RankingPipeline(
_context_feature_columns(),
_example_feature_columns(),
hparams=hparams,
estimator=estimator,
label_feature_name="utility",
label_feature_type=tf.float32)
pip_listwise_export.train_and_eval()
# Does not support non-local training.
with self.assertRaises(ValueError):
pip.train_and_eval(local_training=False)
def build_best_exporter_pipeline(self, hparams):
estimator = tfr.estimator.EstimatorBuilder(
context_feature_columns=_context_feature_columns(),
example_feature_columns=_example_feature_columns(),
scoring_function=self.scoring_function,
hparams=hparams).make_estimator()
return pipeline.RankingPipeline(
context_feature_columns=_context_feature_columns(),
example_feature_columns=_example_feature_columns(),
hparams=hparams,
estimator=estimator,
label_feature_name="utility",
label_feature_type=tf.float32,
best_exporter_metric="metric/ndcg_5")
def build_pipeline(self, hparams, size_feature_name=None):
estimator = tfr.estimator.EstimatorBuilder(
context_feature_columns=_context_feature_columns(),
example_feature_columns=_example_feature_columns(),
scoring_function=self.scoring_function,
hparams=hparams).make_estimator()
return pipeline.RankingPipeline(
context_feature_columns=_context_feature_columns(),
example_feature_columns=_example_feature_columns(),
hparams=hparams,
estimator=estimator,
label_feature_name="utility",
label_feature_type=tf.float32,
size_feature_name=size_feature_name)
def _create_pipeline(self):
hparams = _make_hparams(train_input_pattern=self._data_file,
eval_input_pattern=self._data_file,
model_dir=self._model_dir,
list_size=5)
estimator = tfr.estimator.EstimatorBuilder(
_context_feature_columns(),
_example_feature_columns(),
_scoring_function,
hparams=hparams).make_estimator()
return pipeline.RankingPipeline(_context_feature_columns(),
_example_feature_columns(),
hparams=hparams,
estimator=estimator,
label_feature_name="utility",
label_feature_type=tf.float32)