tfr.keras.pipeline.SimpleDatasetBuilder

Builds datasets from feature specs with a single label spec.

Inherits From: BaseDatasetBuilder, AbstractDatasetBuilder

tfr.keras.pipeline.SimpleDatasetBuilder(
    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]],
    mask_feature_name: str,
    label_spec: Tuple[str, tf.io.FixedLenFeature],
    hparams: tfr.keras.pipeline.DatasetHparams,
    sample_weight_spec: Optional[Tuple[str, tf.io.FixedLenFeature]] = None
)

This supports a single dataset with a single label, which is supposed to be a dense Tensor.

Example usage:

context_feature_spec = {}
example_feature_spec = {
    "example_feature_1": tf.io.FixedLenFeature(
        shape=(1,), dtype=tf.float32, default_value=0.0)
}
mask_feature_name = "list_mask"
label_spec = {
    "utility": tf.io.FixedLenFeature(
        shape=(1,), dtype=tf.float32, default_value=0.0)
}
dataset_hparams = DatasetHparams(
    train_input_pattern="train.dat",
    valid_input_pattern="valid.dat",
    train_batch_size=128,
    valid_batch_size=128)
dataset_builder = SimpleDatasetBuilder(
    context_feature_spec,
    example_feature_spec,
    mask_feature_name,
    label_spec,
    dataset_hparams)

context_feature_spec Maps context (aka, query) names to feature specs.
example_feature_spec Maps example (aka, document) names to feature specs.
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.
label_spec A tuple of the label name and a tf.io.FixedLenFeature spec, or a dict that maps task name to label spec in multi-task setting.
hparams A dict containing model hyperparameters.
sample_weight_spec Feature spec for per-example weight.

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.