tfr.keras.pipeline.BaseDatasetBuilder | TensorFlow Ranking (original) (raw)

Builds datasets from feature specs.

tfr.keras.pipeline.BaseDatasetBuilder(
    context_feature_spec: Dict[str, Union[tf.io.FixedLenFeature, tf.io.VarLenFeature, tf.io.
        RaggedFeature]],
    example_feature_spec: Dict[str, Union[tf.io.FixedLenFeature, tf.io.VarLenFeature, tf.io.
        RaggedFeature]],
    training_only_example_spec: Dict[str, Union[tf.io.FixedLenFeature, tf.io.VarLenFeature, tf.io.
        RaggedFeature]],
    mask_feature_name: str,
    hparams: tfr.keras.pipeline.DatasetHparams,
    training_only_context_spec: Optional[Dict[str, Union[tf.io.FixedLenFeature, tf.io.VarLenFeature, tf.io.
        RaggedFeature]]] = None
)

The BaseDatasetBuilder class is an abstract class inherit fromAbstractDatasetBuilder to serve training and validation datasets and signatures for training ModelFitPipeline.

To be implemented by subclasses:

_features_and_labels(): Contains the logic to map a dict of tensors of dataset to feature tensors and label tensors.

Example subclass implementation:

class SimpleDatasetBuilder(BaseDatasetBuilder):

  def _features_and_labels(self, features):
    label = features.pop("utility")
    return features, label

Args
context_feature_spec	Maps context (aka, query) names to feature specs.
example_feature_spec	Maps example (aka, document) names to feature specs.
training_only_example_spec	Feature specs used for training only like labels and per-example weights.
mask_feature_name	If set, populates the feature dictionary with this name and the coresponding value is a tf.bool Tensor of shape [batch_size, list_size] indicating the actual example is padded or not.
hparams	A dict containing model hyperparameters.
training_only_context_spec	Feature specs used for training only per-list weights.

Methods

`build_signatures`

View source

build_signatures(
    model: tf.keras.Model
) -> Any

See AbstractDatasetBuilder.

`build_train_dataset`

View source

build_train_dataset() -> tf.data.Dataset

See AbstractDatasetBuilder.

`build_valid_dataset`

View source

build_valid_dataset() -> tf.data.Dataset

See AbstractDatasetBuilder.