tfr.keras.pipeline.BaseDatasetBuilder

Builds datasets from feature specs.

Inherits From: AbstractDatasetBuilder

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 from AbstractDatasetBuilder 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

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.